<commit_before>package ledger

import "time"

// TransactionsInDateRange returns a new array of transactions that are in the date range
// specified by start and end. The returned list contains transactions on the same day as start
// but does not include any transactions on the day of end.
func TransactionsInDateRange(trans []*Transaction, start, end time.Time) []*Transaction {
	var newlist []*Transaction

	start = start.Add(-1 * time.Second)

	for _, tran := range trans {
		if tran.Date.After(start) && tran.Date.Before((end)) {
			newlist = append(newlist, tran)
		}
	}

	return newlist
}

// Period is used to specify the length of a date range or frequency
type Period string

// Periods suppored by ledger
const (
	PeriodWeek     Period = "Weekly"
	Period2Week    Period = "BiWeekly"
	PeriodMonth    Period = "Monthly"
	Period2Month   Period = "BiMonthly"
	PeriodQuarter  Period = "Quarterly"
	PeriodSemiYear Period = "SemiYearly"
	PeriodYear     Period = "Yearly"
)

func getDateBoundaries(per Period, start, end time.Time) []time.Time {
	var incDays, incMonth, incYear int
	var periodStart time.Time

	switch per {
	case PeriodWeek:
		incDays = 7
		for periodStart = time.Date(start.Year(), start.Month(), start.Day(), 0, 0, 0, 0, time.UTC); periodStart.Weekday() != time.Sunday; {
			periodStart = periodStart.AddDate(0, 0, -1)
		}
	case Period2Week:
		incDays = 14
		for periodStart = time.Date(start.Year(), start.Month(), start.Day(), 0, 0, 0, 0, time.UTC); periodStart.Weekday() != time.Sunday; {
			periodStart = periodStart.AddDate(0, 0, -1)
		}
	case PeriodMonth:
		incMonth = 1
		periodStart = time.Date(start.Year(), start.Month(), 1, 0, 0, 0, 0, time.UTC)
	case Period2Month:
		incMonth = 2
		periodStart = time.Date(start.Year(), start.Month(), 1, 0, 0, 0, 0, time.UTC)
	case PeriodQuarter:
		incMonth = 3
		switch start.Month() {
		case time.January, time.February, time.March:
			periodStart = time.Date(start.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
		case time.April, time.May, time.June:
			periodStart = time.Date(start.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
		case time.July, time.August, time.September:
			periodStart = time.Date(start.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
		default:
			periodStart = time.Date(start.Year(), time.October, 1, 0, 0, 0, 0, time.UTC)
		}
	case PeriodSemiYear:
		incMonth = 6
		switch start.Month() {
		case time.January, time.February, time.March, time.April, time.May, time.June:
			periodStart = time.Date(start.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
		default:
			periodStart = time.Date(start.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
		}
	case PeriodYear:
		incYear = 1
		periodStart = time.Date(start.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
	}

	boundaries := []time.Time{periodStart}
	for periodStart.Before(end) || periodStart.Equal(end) {
		periodStart = periodStart.AddDate(incYear, incMonth, incDays)
		boundaries = append(boundaries, periodStart)
	}

	return boundaries
}

// RangeType is used to specify how the data is "split" into sections
type RangeType string

const (
	// RangeSnapshot will have each section be the running total at the time of the snapshot
	RangeSnapshot RangeType = "Snapshot"

	// RangePartition will have each section be the accumulated value of the transactions within that partition's date range
	RangePartition RangeType = "Partition"
)

// RangeTransactions contains the transactions and the start and end time of the date range
type RangeTransactions struct {
	Start, End   time.Time
	Transactions []*Transaction
}

// TransactionsByPeriod will return the transactions for each period.
func TransactionsByPeriod(trans []*Transaction, per Period) []*RangeTransactions {
	var results []*RangeTransactions
	if len(trans) < 1 {
		return results
	}

	tStart := trans[0].Date
	tEnd := trans[len(trans)-1].Date

	boundaries := getDateBoundaries(per, tStart, tEnd)

	bStart := boundaries[0]
	for _, boundary := range boundaries[1:] {
		bEnd := boundary

		bTrans := TransactionsInDateRange(trans, bStart, bEnd)
		// End date should be the last day (inclusive, so subtract 1 day)
		results = append(results, &RangeTransactions{Start: bStart, End: bEnd.AddDate(0, 0, -1), Transactions: bTrans})

		bStart = bEnd
	}

	return results
}

// RangeBalance contains the account balances and the start and end time of the date range
type RangeBalance struct {
	Start, End time.Time
	Balances   []*Account
}

// BalancesByPeriod will return the account balances for each period.
func BalancesByPeriod(trans []*Transaction, per Period, rType RangeType) []*RangeBalance {
	var results []*RangeBalance
	if len(trans) < 1 {
		return results
	}

	tStart := trans[0].Date
	tEnd := trans[len(trans)-1].Date

	boundaries := getDateBoundaries(per, tStart, tEnd)

	bStart := boundaries[0]
	for _, boundary := range boundaries[1:] {
		bEnd := boundary

		bTrans := TransactionsInDateRange(trans, bStart, bEnd)
		// End date should be the last day (inclusive, so subtract 1 day)
		results = append(results, &RangeBalance{Start: bStart, End: bEnd.AddDate(0, 0, -1), Balances: GetBalances(bTrans, []string{})})

		if rType == RangePartition {
			bStart = bEnd
		}
	}

	return results
}
<commit_msg>fix partitioning transactions by date and period<commit_after>package ledger

import "time"

// TransactionsInDateRange returns a new array of transactions that are in the date range
// specified by start and end. The returned list contains transactions on the same day as start
// but does not include any transactions on the day of end.
func TransactionsInDateRange(trans []*Transaction, start, end time.Time) []*Transaction {
	var newlist []*Transaction

	start = start.Add(-1 * time.Second)

	for _, tran := range trans {
		if tran.Date.After(start) && tran.Date.Before(end) {
			newlist = append(newlist, tran)
		}
	}

	return newlist
}

// Period is used to specify the length of a date range or frequency
type Period string

// Periods suppored by ledger
const (
	PeriodWeek     Period = "Weekly"
	Period2Week    Period = "BiWeekly"
	PeriodMonth    Period = "Monthly"
	Period2Month   Period = "BiMonthly"
	PeriodQuarter  Period = "Quarterly"
	PeriodSemiYear Period = "SemiYearly"
	PeriodYear     Period = "Yearly"
)

func getDateBoundaries(per Period, start, end time.Time) []time.Time {
	var incDays, incMonth, incYear int
	var periodStart time.Time

	switch per {
	case PeriodWeek:
		incDays = 7
		for periodStart = time.Date(start.Year(), start.Month(), start.Day(), 0, 0, 0, 0, time.UTC); periodStart.Weekday() != time.Sunday; {
			periodStart = periodStart.AddDate(0, 0, -1)
		}
	case Period2Week:
		incDays = 14
		for periodStart = time.Date(start.Year(), start.Month(), start.Day(), 0, 0, 0, 0, time.UTC); periodStart.Weekday() != time.Sunday; {
			periodStart = periodStart.AddDate(0, 0, -1)
		}
	case PeriodMonth:
		incMonth = 1
		periodStart = time.Date(start.Year(), start.Month(), 1, 0, 0, 0, 0, time.UTC)
	case Period2Month:
		incMonth = 2
		periodStart = time.Date(start.Year(), start.Month(), 1, 0, 0, 0, 0, time.UTC)
	case PeriodQuarter:
		incMonth = 3
		switch start.Month() {
		case time.January, time.February, time.March:
			periodStart = time.Date(start.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
		case time.April, time.May, time.June:
			periodStart = time.Date(start.Year(), time.April, 1, 0, 0, 0, 0, time.UTC)
		case time.July, time.August, time.September:
			periodStart = time.Date(start.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
		default:
			periodStart = time.Date(start.Year(), time.October, 1, 0, 0, 0, 0, time.UTC)
		}
	case PeriodSemiYear:
		incMonth = 6
		switch start.Month() {
		case time.January, time.February, time.March, time.April, time.May, time.June:
			periodStart = time.Date(start.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
		default:
			periodStart = time.Date(start.Year(), time.July, 1, 0, 0, 0, 0, time.UTC)
		}
	case PeriodYear:
		incYear = 1
		periodStart = time.Date(start.Year(), time.January, 1, 0, 0, 0, 0, time.UTC)
	}

	boundaries := []time.Time{periodStart}
	for periodStart.Before(end) || periodStart.Equal(end) {
		periodStart = periodStart.AddDate(incYear, incMonth, incDays)
		boundaries = append(boundaries, periodStart)
	}

	return boundaries
}

// RangeType is used to specify how the data is "split" into sections
type RangeType string

const (
	// RangeSnapshot will have each section be the running total at the time of the snapshot
	RangeSnapshot RangeType = "Snapshot"

	// RangePartition will have each section be the accumulated value of the transactions within that partition's date range
	RangePartition RangeType = "Partition"
)

// RangeTransactions contains the transactions and the start and end time of the date range
type RangeTransactions struct {
	Start, End   time.Time
	Transactions []*Transaction
}

// startEndTime will return the start and end Times of a list of transactions
func startEndTime(trans []*Transaction) (start, end time.Time) {
	if len(trans) < 1 {
		return
	}

	start = trans[0].Date
	end = trans[0].Date

	for _, t := range trans {
		if end.Before(t.Date) {
			end = t.Date
		}
		if start.After(t.Date) {
			start = t.Date
		}
	}

	return
}

// TransactionsByPeriod will return the transactions for each period.
func TransactionsByPeriod(trans []*Transaction, per Period) []*RangeTransactions {
	var results []*RangeTransactions
	if len(trans) < 1 {
		return results
	}

	tStart, tEnd := startEndTime(trans)

	boundaries := getDateBoundaries(per, tStart, tEnd)

	bStart := boundaries[0]
	for _, boundary := range boundaries[1:] {
		bEnd := boundary

		bTrans := TransactionsInDateRange(trans, bStart, bEnd)
		// End date should be the last day (inclusive, so subtract 1 day)
		results = append(results, &RangeTransactions{Start: bStart, End: bEnd.AddDate(0, 0, -1), Transactions: bTrans})

		bStart = bEnd
	}

	return results
}

// RangeBalance contains the account balances and the start and end time of the date range
type RangeBalance struct {
	Start, End time.Time
	Balances   []*Account
}

// BalancesByPeriod will return the account balances for each period.
func BalancesByPeriod(trans []*Transaction, per Period, rType RangeType) []*RangeBalance {
	var results []*RangeBalance
	if len(trans) < 1 {
		return results
	}

	tStart, tEnd := startEndTime(trans)

	boundaries := getDateBoundaries(per, tStart, tEnd)

	bStart := boundaries[0]
	for _, boundary := range boundaries[1:] {
		bEnd := boundary

		bTrans := TransactionsInDateRange(trans, bStart, bEnd)
		// End date should be the last day (inclusive, so subtract 1 day)
		results = append(results, &RangeBalance{Start: bStart, End: bEnd.AddDate(0, 0, -1), Balances: GetBalances(bTrans, []string{})})

		if rType == RangePartition {
			bStart = bEnd
		}
	}

	return results
}
<|endoftext|>


<commit_before>// Package logr defines abstract interfaces for logging.  Packages can depend on
// these interfaces and callers can implement logging in whatever way is
// appropriate.
//
// This design derives from Dave Cheney's blog:
//     http://dave.cheney.net/2015/11/05/lets-talk-about-logging
//
// This is a BETA grade API.  Until there is a significant 2nd implementation,
// I don't really know how it will change.
//
// The logging specifically makes it non-trivial to use format strings, to encourage
// attaching structured information instead of unstructured format strings.
//
// Usage
//
// Logging is done using a Logger.  Loggers can have name prefixes and named values
// attached, so that all log messages logged with that Logger have some base context
// associated.
//
// The term "key" is used to refer to the name associated with a particular value, to
// disambiguate it from the general Logger name.
//
// For instance, suppose we're trying to reconcile the state of an object, and we want
// to log that we've made some decision.
//
// With the traditional log package, we might write
//  log.Printf(
//      "decided to set field foo to value %q for object %s/%s",
//       targetValue, object.Namespace, object.Name)
//
// With logr's structured logging, we'd write
//  // elsewhere in the file, set up the logger to log with the prefix of "reconcilers",
//  // and the named value target-type=Foo, for extra context.
//  log := mainLogger.WithName("reconcilers").WithValues("target-type", "Foo")
//
//  // later on...
//  log.Info("setting field foo on object", "value", targetValue, "object", object)
//
// Depending on our logging implementation, we could then make logging decisions based on field values
// (like only logging such events for objects in a certain namespace), or copy the structured
// information into a structured log store.
//
// For logging errors, Logger has a method called Error.  Suppose we wanted to log an
// error while reconciling.  With the traditional log package, we might write
//   log.Errorf("unable to reconcile object %s/%s: %v", object.Namespace, object.Name, err)
//
// With logr, we'd instead write
//   // assuming the above setup for log
//   log.Error(err, "unable to reconcile object", "object", object)
//
// This functions similarly to:
//   log.Info("unable to reconcile object", "error", err, "object", object)
//
// However, it ensures that a standard key for the error value ("error") is used across all
// error logging.  Furthermore, certain implementations may choose to attach additional
// information (such as stack traces) on calls to Error, so it's preferred to use Error
// to log errors.
//
// Parts of a log line
//
// Each log message from a Logger has four types of context:
// logger name, log verbosity, log message, and the named values.
//
// The Logger name constists of a series of name "segments" added by successive calls to WithName.
// These name segments will be joined in some way by the underlying implementation.  It is strongly
// reccomended that name segements contain simple identifiers (letters, digits, and hyphen), and do
// not contain characters that could muddle the log output or confuse the joining operation (e.g.
// whitespace, commas, periods, slashes, brackets, quotes, etc).
//
// Log verbosity represents how little a log matters.  Level zero, the default, matters most.
// Increasing levels matter less and less.  Try to avoid lots of different verbosity levels,
// and instead provide useful keys, logger names, and log messages for users to filter on.
//
// The log message consists of a constant message attached to the the log line.  This
// should generally be a simple description of what's occuring, and should never be a format string.
//
// Variable information can then be attached using named values (key/value pairs).  Keys are arbitrary
// strings, while values may be any Go value.
//
// Key Naming Conventions
//
// While users are generally free to use key names of their choice, it's generally best to avoid
// using the following keys, as they're frequently used by implementations:
//
// - `"error"`: the underlying error value in the `Error` method.
// - `"stacktrace"`: the stack trace associated with a particular log line or error
//                   (often from the `Error` message).
// - `"caller"`: the calling information (file/line) of a particular log line.
// - `"msg"`: the log message.
// - `"level"`: the log level.
// - `"ts"`: the timestamp for a log line.
//
// Implementations are encouraged to make use of these keys to represent the above
// concepts, when neccessary (for example, in a pure-JSON output form, it would be
// necessary to represent at least message and timestamp as ordinary named values).
package logr

// TODO: consider adding back in format strings if they're really needed
// TODO: consider other bits of zap/zapcore functionality like ObjectMarshaller (for arbitrary objects)
// TODO: consider other bits of glog functionality like Flush, InfoDepth, OutputStats

// InfoLogger represents the ability to log non-error messages, at a particular verbosity.
type InfoLogger interface {
	// Info logs a non-error message with the given key/value pairs as context.
	//
	// The msg argument should be used to add some constant description to
	// the log line.  The key/value pairs can then be used to add additional
	// variable information.  The key/value pairs should alternate string
	// keys and arbitrary values.
	Info(msg string, keysAndValues ...interface{})

	// Enabled tests whether this InfoLogger is enabled.  For example,
	// commandline flags might be used to set the logging verbosity and disable
	// some info logs.
	Enabled() bool
}

// Logger represents the ability to log messages, both errors and not.
type Logger interface {
	// All Loggers implement InfoLogger.  Calling InfoLogger methods directly on
	// a Logger value is equivalent to calling them on a V(0) InfoLogger.  For
	// example, logger.Info() produces the same result as logger.V(0).Info.
	InfoLogger

	// Error logs an error, with the given message and key/value pairs as context.
	// It functions similarly to calling Info with the "error" named value, but may
	// have unique behavior, and should be preferred for logging errors (see the
	// package documentations for more information).
	//
	// The msg field should be used to add context to any underlying error,
	// while the err field should be used to attach the actual error that
	// triggered this log line, if present.
	Error(err error, msg string, keysAndValues ...interface{})

	// V returns an InfoLogger value for a specific verbosity level.  A higher
	// verbosity level means a log message is less important.
	V(level int) InfoLogger

	// WithValues adds some key-value pairs of context to a logger.
	// See Info for documentation on how key/value pairs work.
	WithValues(keysAndValues ...interface{}) Logger

	// WithName adds a new element to the logger's name.
	// Successive calls with WithName continue to append
	// suffixes to the logger's name.  It's strongly reccomended
	// that name segments contain only letters, digits, and hyphens
	// (see the package documentation for more information).
	WithName(name string) Logger
}
<commit_msg>Forbid log levels less that zero<commit_after>// Package logr defines abstract interfaces for logging.  Packages can depend on
// these interfaces and callers can implement logging in whatever way is
// appropriate.
//
// This design derives from Dave Cheney's blog:
//     http://dave.cheney.net/2015/11/05/lets-talk-about-logging
//
// This is a BETA grade API.  Until there is a significant 2nd implementation,
// I don't really know how it will change.
//
// The logging specifically makes it non-trivial to use format strings, to encourage
// attaching structured information instead of unstructured format strings.
//
// Usage
//
// Logging is done using a Logger.  Loggers can have name prefixes and named values
// attached, so that all log messages logged with that Logger have some base context
// associated.
//
// The term "key" is used to refer to the name associated with a particular value, to
// disambiguate it from the general Logger name.
//
// For instance, suppose we're trying to reconcile the state of an object, and we want
// to log that we've made some decision.
//
// With the traditional log package, we might write
//  log.Printf(
//      "decided to set field foo to value %q for object %s/%s",
//       targetValue, object.Namespace, object.Name)
//
// With logr's structured logging, we'd write
//  // elsewhere in the file, set up the logger to log with the prefix of "reconcilers",
//  // and the named value target-type=Foo, for extra context.
//  log := mainLogger.WithName("reconcilers").WithValues("target-type", "Foo")
//
//  // later on...
//  log.Info("setting field foo on object", "value", targetValue, "object", object)
//
// Depending on our logging implementation, we could then make logging decisions based on field values
// (like only logging such events for objects in a certain namespace), or copy the structured
// information into a structured log store.
//
// For logging errors, Logger has a method called Error.  Suppose we wanted to log an
// error while reconciling.  With the traditional log package, we might write
//   log.Errorf("unable to reconcile object %s/%s: %v", object.Namespace, object.Name, err)
//
// With logr, we'd instead write
//   // assuming the above setup for log
//   log.Error(err, "unable to reconcile object", "object", object)
//
// This functions similarly to:
//   log.Info("unable to reconcile object", "error", err, "object", object)
//
// However, it ensures that a standard key for the error value ("error") is used across all
// error logging.  Furthermore, certain implementations may choose to attach additional
// information (such as stack traces) on calls to Error, so it's preferred to use Error
// to log errors.
//
// Parts of a log line
//
// Each log message from a Logger has four types of context:
// logger name, log verbosity, log message, and the named values.
//
// The Logger name constists of a series of name "segments" added by successive calls to WithName.
// These name segments will be joined in some way by the underlying implementation.  It is strongly
// reccomended that name segements contain simple identifiers (letters, digits, and hyphen), and do
// not contain characters that could muddle the log output or confuse the joining operation (e.g.
// whitespace, commas, periods, slashes, brackets, quotes, etc).
//
// Log verbosity represents how little a log matters.  Level zero, the default, matters most.
// Increasing levels matter less and less.  Try to avoid lots of different verbosity levels,
// and instead provide useful keys, logger names, and log messages for users to filter on.
// It's illegal to pass a log level below zero.
//
// The log message consists of a constant message attached to the the log line.  This
// should generally be a simple description of what's occuring, and should never be a format string.
//
// Variable information can then be attached using named values (key/value pairs).  Keys are arbitrary
// strings, while values may be any Go value.
//
// Key Naming Conventions
//
// While users are generally free to use key names of their choice, it's generally best to avoid
// using the following keys, as they're frequently used by implementations:
//
// - `"error"`: the underlying error value in the `Error` method.
// - `"stacktrace"`: the stack trace associated with a particular log line or error
//                   (often from the `Error` message).
// - `"caller"`: the calling information (file/line) of a particular log line.
// - `"msg"`: the log message.
// - `"level"`: the log level.
// - `"ts"`: the timestamp for a log line.
//
// Implementations are encouraged to make use of these keys to represent the above
// concepts, when neccessary (for example, in a pure-JSON output form, it would be
// necessary to represent at least message and timestamp as ordinary named values).
package logr

// TODO: consider adding back in format strings if they're really needed
// TODO: consider other bits of zap/zapcore functionality like ObjectMarshaller (for arbitrary objects)
// TODO: consider other bits of glog functionality like Flush, InfoDepth, OutputStats

// InfoLogger represents the ability to log non-error messages, at a particular verbosity.
type InfoLogger interface {
	// Info logs a non-error message with the given key/value pairs as context.
	//
	// The msg argument should be used to add some constant description to
	// the log line.  The key/value pairs can then be used to add additional
	// variable information.  The key/value pairs should alternate string
	// keys and arbitrary values.
	Info(msg string, keysAndValues ...interface{})

	// Enabled tests whether this InfoLogger is enabled.  For example,
	// commandline flags might be used to set the logging verbosity and disable
	// some info logs.
	Enabled() bool
}

// Logger represents the ability to log messages, both errors and not.
type Logger interface {
	// All Loggers implement InfoLogger.  Calling InfoLogger methods directly on
	// a Logger value is equivalent to calling them on a V(0) InfoLogger.  For
	// example, logger.Info() produces the same result as logger.V(0).Info.
	InfoLogger

	// Error logs an error, with the given message and key/value pairs as context.
	// It functions similarly to calling Info with the "error" named value, but may
	// have unique behavior, and should be preferred for logging errors (see the
	// package documentations for more information).
	//
	// The msg field should be used to add context to any underlying error,
	// while the err field should be used to attach the actual error that
	// triggered this log line, if present.
	Error(err error, msg string, keysAndValues ...interface{})

	// V returns an InfoLogger value for a specific verbosity level.  A higher
	// verbosity level means a log message is less important.  It's illegal to
	// pass a log level less than zero.
	V(level int) InfoLogger

	// WithValues adds some key-value pairs of context to a logger.
	// See Info for documentation on how key/value pairs work.
	WithValues(keysAndValues ...interface{}) Logger

	// WithName adds a new element to the logger's name.
	// Successive calls with WithName continue to append
	// suffixes to the logger's name.  It's strongly reccomended
	// that name segments contain only letters, digits, and hyphens
	// (see the package documentation for more information).
	WithName(name string) Logger
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"math/big"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

func main() {
	var fromStr string
	flag.StringVar(&fromStr, "from", "0-1-1", "from date")
	var toStr string
	flag.StringVar(&toStr, "to", "9999-1-1", "to date")
	var cmdProperties bool
	flag.BoolVar(&cmdProperties, "props", false, "show properties")
	var cmdMonthlyExpenses bool
	flag.BoolVar(&cmdMonthlyExpenses, "monthly", false, "show monthly expenses")

	flag.Parse()

	// options
	parseDate := func(str string) time.Time {
		parts := dateSepPattern.Split(str, -1)
		if len(parts) != 3 {
			panic(me(nil, "bad date: %s", str))
		}
		year, err := strconv.Atoi(parts[0])
		ce(err, "parse year: %s", str)
		month, err := strconv.Atoi(parts[1])
		ce(err, "parse month: %s", str)
		day, err := strconv.Atoi(parts[2])
		ce(err, "parse day: %s", str)
		return time.Date(year, time.Month(month), day, 0, 0, 0, 0, time.Local)
	}
	fromTime := parseDate(fromStr).Add(-time.Hour)
	toTime := parseDate(toStr).Add(time.Hour)

	// usage
	args := flag.Args()
	if len(args) < 1 {
		pt("usage: %s [options] <file path>\n", os.Args[0])
		flag.Usage()
		return
	}

	// read ledger file
	ledgerPath := args[0]
	contentBytes, err := ioutil.ReadFile(ledgerPath)
	ce(err, "read ledger")
	content := string(contentBytes)
	content = strings.Replace(content, "\r\n", "\n", -1)
	content = strings.Replace(content, "\r", "\n", -1)

	// parse blocks
	var blocks [][]string
	var block []string
	for _, line := range strings.Split(content, "\n") {
		line = strings.TrimSpace(line)
		if len(line) == 0 {
			if len(block) > 0 {
				blocks = append(blocks, block)
				block = []string{}
			}
		} else {
			block = append(block, line)
		}
	}
	if len(block) > 0 {
		blocks = append(blocks, block)
	}

	// transaction
	type Account struct {
		Name        string
		Subs        map[string]*Account
		Parent      *Account
		Balances    map[string]*big.Rat
		Proportions map[string]*big.Rat
	}
	type Entry struct {
		Account     *Account
		Currency    string
		Amount      *big.Rat
		Description string
	}
	type Transaction struct {
		Year        int
		Month       int
		Day         int
		Time        time.Time
		Description string
		Entries     []*Entry
	}
	var transactions []*Transaction

	rootAccount := &Account{
		Name:        "root",
		Subs:        make(map[string]*Account),
		Parent:      nil,
		Balances:    make(map[string]*big.Rat),
		Proportions: make(map[string]*big.Rat),
	}
	var getAccount func(root *Account, path []string) *Account
	getAccount = func(root *Account, path []string) *Account {
		if len(path) == 0 {
			panic(me(nil, "bad account: %v", path))
		} else if len(path) == 1 {
			name := path[0]
			account, ok := root.Subs[name]
			if ok {
				return account
			}
			account = &Account{
				Name:        name,
				Subs:        make(map[string]*Account),
				Parent:      root,
				Balances:    make(map[string]*big.Rat),
				Proportions: make(map[string]*big.Rat),
			}
			root.Subs[name] = account
			return account
		}
		name := path[0]
		account, ok := root.Subs[name]
		if !ok {
			account = &Account{
				Name:        name,
				Subs:        make(map[string]*Account),
				Parent:      root,
				Balances:    make(map[string]*big.Rat),
				Proportions: make(map[string]*big.Rat),
			}
			root.Subs[name] = account
		}
		return getAccount(account, path[1:])
	}

	// collect transactions
	for _, block := range blocks {
		n := 0
		transaction := new(Transaction)

		// parse
		for _, line := range block {
			if strings.HasPrefix(line, "#") {
				continue
			}
			n++

			if n == 1 {
				// transaction header
				parts := blanksPattern.Split(line, 2)
				if len(parts) != 2 {
					panic(me(nil, "bad transaction header: %s", line))
				}

				dateStr := parts[0]
				dateParts := dateSepPattern.Split(dateStr, -1)
				if len(dateParts) != 3 {
					panic(me(nil, "bad date: %s", line))
				}
				year, err := strconv.Atoi(dateParts[0])
				ce(err, "parse year: %s", line)
				transaction.Year = year
				month, err := strconv.Atoi(dateParts[1])
				ce(err, "parse month: %s", line)
				transaction.Month = month
				day, err := strconv.Atoi(dateParts[2])
				ce(err, "parse day: %s", line)
				transaction.Day = day
				transaction.Time = time.Date(year, time.Month(month), day, 0, 0, 0, 0, time.Local)

				transaction.Description = parts[1]

			} else {
				// entry
				parts := blanksPattern.Split(line, 3)
				if len(parts) < 2 {
					panic(me(nil, "bad entry: %s", line))
				}
				entry := new(Entry)

				accountStr := parts[0]
				account := getAccount(rootAccount, strings.Split(accountStr, "："))
				entry.Account = account

				currency, runeSize := utf8.DecodeRuneInString(parts[1])
				entry.Currency = string(currency)
				amountStr := parts[1][runeSize:]
				amount := new(big.Rat)
				_, ok := amount.SetString(amountStr)
				if !ok {
					panic(me(nil, "bad amount: %s", amountStr))
				}
				entry.Amount = amount

				if len(parts) > 2 {
					entry.Description = parts[2]
				}

				transaction.Entries = append(transaction.Entries, entry)
			}

		}

		if transaction.Year == 0 {
			// empty
			continue
		}
		t := time.Date(transaction.Year, time.Month(transaction.Month), transaction.Day,
			0, 0, 0, 0, time.Local)
		if t.Before(fromTime) || t.After(toTime) {
			// out of range
			continue
		}

		// check balance
		sum := big.NewRat(0, 1)
		for _, entry := range transaction.Entries {
			sum.Add(sum, entry.Amount)
			// update account balance
			account := entry.Account
			for account != nil {
				balance, ok := account.Balances[entry.Currency]
				if !ok {
					balance = big.NewRat(0, 1)
					account.Balances[entry.Currency] = balance
				}
				balance.Add(balance, entry.Amount)
				account = account.Parent
			}
		}
		if !(sum.Cmp(zeroRat) == 0) {
			panic(me(nil, "not balanced: %s", strings.Join(block, "\n")))
		}

		transactions = append(transactions, transaction)
	}

	// calculate proportions
	var calculateProportion func(*Account)
	calculateProportion = func(account *Account) {
		for _, sub := range account.Subs {
			for currency, balance := range sub.Balances {
				if account.Balances[currency].Sign() != 0 {
					b := big.NewRat(0, 1)
					b.Set(balance)
					sub.Proportions[currency] = b.Quo(balance, account.Balances[currency])
					b.Abs(b)
				}
				calculateProportion(sub)
			}
		}
	}
	calculateProportion(rootAccount)

	// print accounts
	onePercent := big.NewRat(1, 100)
	var printAccount func(account *Account, level int)
	printAccount = func(account *Account, level int) {
		allZero := true
		for _, balance := range account.Balances {
			if balance.Cmp(zeroRat) != 0 {
				allZero = false
				break
			}
		}
		if allZero && account != rootAccount {
			return
		}
		pt("%s%s", strings.Repeat(" │    ", level), account.Name)
		var currencyNames []string
		for name := range account.Balances {
			currencyNames = append(currencyNames, name)
		}
		sort.Strings(currencyNames)
		for _, name := range currencyNames {
			balance := account.Balances[name]
			var proportion string
			if p, ok := account.Proportions[name]; ok && p.Cmp(onePercent) > 0 {
				proportion = " " + p.Mul(p, big.NewRat(100, 1)).FloatString(3) + "%"
			}
			pt(" %s%s%s", name, balance.FloatString(2), proportion)
		}
		pt("\n")

		var subNames []string
		for name := range account.Subs {
			subNames = append(subNames, name)
		}
		digitPattern := regexp.MustCompile(`^[0-9]+$`)
		sort.Slice(subNames, func(i, j int) bool {
			a := account.Subs[subNames[i]]
			b := account.Subs[subNames[j]]
			if len(a.Subs) == 0 &&
				len(b.Subs) == 0 &&
				!digitPattern.MatchString(subNames[i]) && // 月份
				!digitPattern.MatchString(subNames[j]) {
				sumA := big.NewRat(0, 1)
				for _, balance := range a.Balances {
					sumA.Add(sumA, balance)
				}
				sumB := big.NewRat(0, 1)
				for _, balance := range b.Balances {
					sumB.Add(sumB, balance)
				}
				return sumA.Cmp(sumB) > 0
			}
			return subNames[i] < subNames[j]
		})
		for _, name := range subNames {
			printAccount(account.Subs[name], level+1)
		}
	}
	printAccount(rootAccount, 0)

	if cmdProperties {
		accountNames := map[string]bool{
			"书籍":   true,
			"保健品":  true,
			"性用品":  true,
			"消耗品":  true,
			"物品":   true,
			"电子":   true,
			"药物":   true,
			"衣物服饰": true,
		}
		accounts := map[*Account]bool{}
		for name, account := range rootAccount.Subs["支出"].Subs {
			if accountNames[name] {
				accounts[account] = true
			}
		}
		var ts []*Transaction
	loop_t:
		for _, t := range transactions {
			for _, entry := range t.Entries {
				if accounts[entry.Account] {
					ts = append(ts, t)
					continue loop_t
				}
			}
		}
		sort.Slice(ts, func(i, j int) bool {
			return ts[i].Time.Before(ts[j].Time)
		})
		for _, t := range ts {
			pt("%s %s\n", t.Time.Format("01-02"), t.Description)
		}
	}

	if cmdMonthlyExpenses {
		accounts := make(map[*Account]bool)
		for _, account := range rootAccount.Subs["支出"].Subs {
			accounts[account] = true
		}
		monthEntries := make(map[string][]*Entry)
		for _, transaction := range transactions {
			for _, entry := range transaction.Entries {
				if accounts[entry.Account] {
					// is expense
					monthStr := fmt.Sprintf("%04d-%02d", transaction.Year, transaction.Month)
					monthEntries[monthStr] = append(monthEntries[monthStr], entry)
				}
			}
		}
		var months []string
		for month := range monthEntries {
			months = append(months, month)
		}
		sort.Strings(months)
		for _, month := range months {
			entries := monthEntries[month]
			sums := make(map[string]*big.Rat)
			for _, entry := range entries {
				if sums[entry.Currency] == nil {
					sums[entry.Currency] = big.NewRat(0, 1)
				}
				sums[entry.Currency].Add(sums[entry.Currency], entry.Amount)
			}
			pt("%s", month)
			for cur, sum := range sums {
				pt(" %s%s", cur, sum.FloatString(2))
			}
			pt("\n")
		}
	}

}
<commit_msg>update<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"math/big"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"
)

func main() {
	var fromStr string
	flag.StringVar(&fromStr, "from", "0-1-1", "from date")
	var toStr string
	flag.StringVar(&toStr, "to", "9999-1-1", "to date")
	var cmdProperties bool
	flag.BoolVar(&cmdProperties, "props", false, "show properties")
	var cmdMonthlyExpenses bool
	flag.BoolVar(&cmdMonthlyExpenses, "monthly", false, "show monthly expenses")

	flag.Parse()

	// options
	parseDate := func(str string) time.Time {
		parts := dateSepPattern.Split(str, -1)
		if len(parts) != 3 {
			panic(me(nil, "bad date: %s", str))
		}
		year, err := strconv.Atoi(parts[0])
		ce(err, "parse year: %s", str)
		month, err := strconv.Atoi(parts[1])
		ce(err, "parse month: %s", str)
		day, err := strconv.Atoi(parts[2])
		ce(err, "parse day: %s", str)
		return time.Date(year, time.Month(month), day, 0, 0, 0, 0, time.Local)
	}
	fromTime := parseDate(fromStr).Add(-time.Hour)
	toTime := parseDate(toStr).Add(time.Hour)

	// usage
	args := flag.Args()
	if len(args) < 1 {
		pt("usage: %s [options] <file path>\n", os.Args[0])
		flag.Usage()
		return
	}

	// read ledger file
	ledgerPath := args[0]
	contentBytes, err := ioutil.ReadFile(ledgerPath)
	ce(err, "read ledger")
	content := string(contentBytes)
	content = strings.Replace(content, "\r\n", "\n", -1)
	content = strings.Replace(content, "\r", "\n", -1)

	// parse blocks
	var blocks [][]string
	var block []string
	for _, line := range strings.Split(content, "\n") {
		line = strings.TrimSpace(line)
		if len(line) == 0 {
			if len(block) > 0 {
				blocks = append(blocks, block)
				block = []string{}
			}
		} else {
			block = append(block, line)
		}
	}
	if len(block) > 0 {
		blocks = append(blocks, block)
	}

	// transaction
	type Account struct {
		Name        string
		Subs        map[string]*Account
		Parent      *Account
		Balances    map[string]*big.Rat
		Proportions map[string]*big.Rat
	}
	type Entry struct {
		Account     *Account
		Currency    string
		Amount      *big.Rat
		Description string
	}
	type Transaction struct {
		Year        int
		Month       int
		Day         int
		Time        time.Time
		Description string
		Entries     []*Entry
	}
	var transactions []*Transaction

	rootAccount := &Account{
		Name:        "root",
		Subs:        make(map[string]*Account),
		Parent:      nil,
		Balances:    make(map[string]*big.Rat),
		Proportions: make(map[string]*big.Rat),
	}
	var getAccount func(root *Account, path []string) *Account
	getAccount = func(root *Account, path []string) *Account {
		if len(path) == 0 {
			panic(me(nil, "bad account: %v", path))
		} else if len(path) == 1 {
			name := path[0]
			account, ok := root.Subs[name]
			if ok {
				return account
			}
			account = &Account{
				Name:        name,
				Subs:        make(map[string]*Account),
				Parent:      root,
				Balances:    make(map[string]*big.Rat),
				Proportions: make(map[string]*big.Rat),
			}
			root.Subs[name] = account
			return account
		}
		name := path[0]
		account, ok := root.Subs[name]
		if !ok {
			account = &Account{
				Name:        name,
				Subs:        make(map[string]*Account),
				Parent:      root,
				Balances:    make(map[string]*big.Rat),
				Proportions: make(map[string]*big.Rat),
			}
			root.Subs[name] = account
		}
		return getAccount(account, path[1:])
	}

	// collect transactions
	for _, block := range blocks {
		n := 0
		transaction := new(Transaction)

		// parse
		for _, line := range block {
			if strings.HasPrefix(line, "#") {
				continue
			}
			n++

			if n == 1 {
				// transaction header
				parts := blanksPattern.Split(line, 2)
				if len(parts) != 2 {
					panic(me(nil, "bad transaction header: %s", line))
				}

				dateStr := parts[0]
				dateParts := dateSepPattern.Split(dateStr, -1)
				if len(dateParts) != 3 {
					panic(me(nil, "bad date: %s", line))
				}
				year, err := strconv.Atoi(dateParts[0])
				ce(err, "parse year: %s", line)
				transaction.Year = year
				month, err := strconv.Atoi(dateParts[1])
				ce(err, "parse month: %s", line)
				transaction.Month = month
				day, err := strconv.Atoi(dateParts[2])
				ce(err, "parse day: %s", line)
				transaction.Day = day
				transaction.Time = time.Date(year, time.Month(month), day, 0, 0, 0, 0, time.Local)

				transaction.Description = parts[1]

			} else {
				// entry
				parts := blanksPattern.Split(line, 3)
				if len(parts) < 2 {
					panic(me(nil, "bad entry: %s", line))
				}
				entry := new(Entry)

				accountStr := parts[0]
				account := getAccount(rootAccount, strings.Split(accountStr, "："))
				entry.Account = account

				currency, runeSize := utf8.DecodeRuneInString(parts[1])
				entry.Currency = string(currency)
				amountStr := parts[1][runeSize:]
				amount := new(big.Rat)
				_, ok := amount.SetString(amountStr)
				if !ok {
					panic(me(nil, "bad amount: %s", amountStr))
				}
				entry.Amount = amount

				if len(parts) > 2 {
					entry.Description = parts[2]
				}

				transaction.Entries = append(transaction.Entries, entry)
			}

		}

		if transaction.Year == 0 {
			// empty
			continue
		}
		t := time.Date(transaction.Year, time.Month(transaction.Month), transaction.Day,
			0, 0, 0, 0, time.Local)
		if t.Before(fromTime) || t.After(toTime) {
			// out of range
			continue
		}

		// check balance
		sum := big.NewRat(0, 1)
		for _, entry := range transaction.Entries {
			sum.Add(sum, entry.Amount)
			// update account balance
			account := entry.Account
			for account != nil {
				balance, ok := account.Balances[entry.Currency]
				if !ok {
					balance = big.NewRat(0, 1)
					account.Balances[entry.Currency] = balance
				}
				balance.Add(balance, entry.Amount)
				account = account.Parent
			}
		}
		if !(sum.Cmp(zeroRat) == 0) {
			panic(me(nil, "not balanced: %s", strings.Join(block, "\n")))
		}

		transactions = append(transactions, transaction)
	}

	// calculate proportions
	var calculateProportion func(*Account)
	calculateProportion = func(account *Account) {
		for _, sub := range account.Subs {
			for currency, balance := range sub.Balances {
				if account.Balances[currency].Sign() != 0 {
					b := big.NewRat(0, 1)
					b.Set(balance)
					sub.Proportions[currency] = b.Quo(balance, account.Balances[currency])
					b.Abs(b)
				}
				calculateProportion(sub)
			}
		}
	}
	calculateProportion(rootAccount)

	// print accounts
	var printAccount func(account *Account, level int)
	printAccount = func(account *Account, level int) {
		allZero := true
		for _, balance := range account.Balances {
			if balance.Cmp(zeroRat) != 0 {
				allZero = false
				break
			}
		}
		if allZero && account != rootAccount {
			return
		}
		pt("%s%s", strings.Repeat(" │    ", level), account.Name)
		var currencyNames []string
		for name := range account.Balances {
			currencyNames = append(currencyNames, name)
		}
		sort.Strings(currencyNames)
		for _, name := range currencyNames {
			balance := account.Balances[name]
			var proportion string
			if p, ok := account.Proportions[name]; ok {
				proportion = " " + p.Mul(p, big.NewRat(100, 1)).FloatString(3) + "%"
			}
			pt(" %s%s%s", name, balance.FloatString(2), proportion)
		}
		pt("\n")

		var subNames []string
		for name := range account.Subs {
			subNames = append(subNames, name)
		}
		digitPattern := regexp.MustCompile(`^[0-9]+$`)
		sort.Slice(subNames, func(i, j int) bool {
			a := account.Subs[subNames[i]]
			b := account.Subs[subNames[j]]
			if len(a.Subs) == 0 &&
				len(b.Subs) == 0 &&
				!digitPattern.MatchString(subNames[i]) && // 月份
				!digitPattern.MatchString(subNames[j]) {
				sumA := big.NewRat(0, 1)
				for _, balance := range a.Balances {
					sumA.Add(sumA, balance)
				}
				sumB := big.NewRat(0, 1)
				for _, balance := range b.Balances {
					sumB.Add(sumB, balance)
				}
				return sumA.Cmp(sumB) > 0
			}
			return subNames[i] < subNames[j]
		})
		for _, name := range subNames {
			printAccount(account.Subs[name], level+1)
		}
	}
	printAccount(rootAccount, 0)

	if cmdProperties {
		accountNames := map[string]bool{
			"书籍":   true,
			"保健品":  true,
			"性用品":  true,
			"消耗品":  true,
			"物品":   true,
			"电子":   true,
			"药物":   true,
			"衣物服饰": true,
		}
		accounts := map[*Account]bool{}
		for name, account := range rootAccount.Subs["支出"].Subs {
			if accountNames[name] {
				accounts[account] = true
			}
		}
		var ts []*Transaction
	loop_t:
		for _, t := range transactions {
			for _, entry := range t.Entries {
				if accounts[entry.Account] {
					ts = append(ts, t)
					continue loop_t
				}
			}
		}
		sort.Slice(ts, func(i, j int) bool {
			return ts[i].Time.Before(ts[j].Time)
		})
		for _, t := range ts {
			pt("%s %s\n", t.Time.Format("01-02"), t.Description)
		}
	}

	if cmdMonthlyExpenses {
		accounts := make(map[*Account]bool)
		for _, account := range rootAccount.Subs["支出"].Subs {
			accounts[account] = true
		}
		monthEntries := make(map[string][]*Entry)
		for _, transaction := range transactions {
			for _, entry := range transaction.Entries {
				if accounts[entry.Account] {
					// is expense
					monthStr := fmt.Sprintf("%04d-%02d", transaction.Year, transaction.Month)
					monthEntries[monthStr] = append(monthEntries[monthStr], entry)
				}
			}
		}
		var months []string
		for month := range monthEntries {
			months = append(months, month)
		}
		sort.Strings(months)
		for _, month := range months {
			entries := monthEntries[month]
			sums := make(map[string]*big.Rat)
			for _, entry := range entries {
				if sums[entry.Currency] == nil {
					sums[entry.Currency] = big.NewRat(0, 1)
				}
				sums[entry.Currency].Add(sums[entry.Currency], entry.Amount)
			}
			pt("%s", month)
			for cur, sum := range sums {
				pt(" %s%s", cur, sum.FloatString(2))
			}
			pt("\n")
		}
	}

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/gorilla/feeds"
	"github.com/joho/godotenv"
	"net/http"
	"os"
)

func RenderFeed(w http.ResponseWriter, feed *feeds.Feed) {
	atom, err := feed.ToAtom()
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/atom+xml")
	fmt.Fprintln(w, atom)
}

func HandlePurolandNews(w http.ResponseWriter, r *http.Request) {
	feed, err := GetPurolandNews()
	if err != nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	RenderFeed(w, feed)
}

func HandlePurolandInfo(w http.ResponseWriter, r *http.Request) {
	feed, err := GetPurolandInfo()
	if err != nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	RenderFeed(w, feed)
}

func HandleCharacterShow(w http.ResponseWriter, r *http.Request) {
	feed, err := GetCharacterShow()
	if err != nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	RenderFeed(w, feed)
}

func main() {
	_ = godotenv.Load()

	http.HandleFunc("/puroland-info", HandlePurolandInfo)
	http.HandleFunc("/puroland-news", HandlePurolandNews)
	http.HandleFunc("/character-show", HandleCharacterShow)

	port := os.Getenv("PORT")
	if port == "" {
		port = "13000"
	}
	http.ListenAndServe(":" + port, nil)
}
<commit_msg>go fmt<commit_after>package main

import (
	"fmt"
	"github.com/gorilla/feeds"
	"github.com/joho/godotenv"
	"net/http"
	"os"
)

func RenderFeed(w http.ResponseWriter, feed *feeds.Feed) {
	atom, err := feed.ToAtom()
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/atom+xml")
	fmt.Fprintln(w, atom)
}

func HandlePurolandNews(w http.ResponseWriter, r *http.Request) {
	feed, err := GetPurolandNews()
	if err != nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	RenderFeed(w, feed)
}

func HandlePurolandInfo(w http.ResponseWriter, r *http.Request) {
	feed, err := GetPurolandInfo()
	if err != nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	RenderFeed(w, feed)
}

func HandleCharacterShow(w http.ResponseWriter, r *http.Request) {
	feed, err := GetCharacterShow()
	if err != nil {
		w.WriteHeader(http.StatusServiceUnavailable)
		return
	}
	RenderFeed(w, feed)
}

func main() {
	_ = godotenv.Load()

	http.HandleFunc("/puroland-info", HandlePurolandInfo)
	http.HandleFunc("/puroland-news", HandlePurolandNews)
	http.HandleFunc("/character-show", HandleCharacterShow)

	port := os.Getenv("PORT")
	if port == "" {
		port = "13000"
	}
	http.ListenAndServe(":"+port, nil)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"regexp"
	"sort"
	"time"

	"github.com/sirupsen/logrus"
	humanize "github.com/dustin/go-humanize"
	"github.com/nlopes/slack"
)

const extractMsgGroupName = "msg"

// extract suffixes of Slack messages starting with @bot-name
var extractMsgPattern = `(?m)^\s*<@%s>:?\s*(?P<` + extractMsgGroupName + `>.*)$`

// A standupMsg is a Slack message directed to the arriba bot (i.e. with a @botname prefix)
type standupMsg struct {
	ts   time.Time
	text string
}

// channelStandup contains the latest standup message of each user in a Slack channel.
type channelStandup map[string]standupMsg

// sortableChannelStandup is a channelStandup, sortable by the timestamp of its messages
// sortableChannelStandup implements sort.Interface to sort the keys of the channelStandup
type sortableChannelStandup struct {
	keys []string
	cs   channelStandup
}

func (s sortableChannelStandup) Swap(i, j int) { s.keys[i], s.keys[j] = s.keys[j], s.keys[i] }
func (s sortableChannelStandup) Len() int      { return len(s.keys) }
func (s sortableChannelStandup) Less(i, j int) bool {
	return s.cs[s.keys[i]].ts.After(s.cs[s.keys[j]].ts)
}

// getKeysByTimestamp returns the userIDs of the standup ordered by their message timestamp (newer first).
func (cs channelStandup) getKeysByTimestamp() []string {
	keys := make([]string, 0, len(cs))
	for k := range cs {
		keys = append(keys, k)
	}
	scs := sortableChannelStandup{
		cs:   cs,
		keys: keys,
	}
	sort.Sort(scs)
	return scs.keys
}

// standups contains the channelStandup of all Slack channels known to the bot.
type standups map[string]channelStandup

// conversation is a generic way to access the IDs, Names and history of both
// slack.Channel and slack.Group. Unfortunately nlopes/slack doesn't expose the
// underlying common type (groupConversation) and we cannot define methods for
// non-local types, which would allow to make things much cleaner ...
type conversation interface {
	getID() string
	getName() string
	getHistory(*slack.RTM, slack.HistoryParameters) (*slack.History, error)
}

type channel slack.Channel

func (c channel) getID() string   { return c.ID }
func (c channel) getName() string { return c.Name }
func (c channel) getHistory(rtm *slack.RTM, params slack.HistoryParameters) (*slack.History, error) {
	return rtm.GetChannelHistory(c.getID(), params)
}

type group slack.Group

func (g group) getID() string   { return g.ID }
func (g group) getName() string { return g.Name }
func (g group) getHistory(rtm *slack.RTM, params slack.HistoryParameters) (*slack.History, error) {
	return rtm.GetGroupHistory(g.getID(), params)
}

type arriba struct {
	rtm              *slack.RTM
	botID            string
	botName          string
	extractMsgRE     *regexp.Regexp
	historyDaysLimit int
	standups         standups
}

func newArriba(rtm *slack.RTM, historyDaysLimit int) arriba {
	return arriba{
		rtm:              rtm,
		historyDaysLimit: historyDaysLimit,
		standups:         make(standups),
	}
}

func parseSlackTimeStamp(ts string) (time.Time, error) {
	var seconds, milliseconds int64
	_, err := fmt.Sscanf(ts, "%d.%d", &seconds, &milliseconds)
	if err != nil {
		logrus.Warn("Can't parse timestamp ", ts)
		return time.Now(), err
	}
	return time.Unix(seconds, milliseconds*1000), nil
}

// extractStandupMsg parses Slack messages starting with @bot-name
func (a arriba) extractChannelStandupMsg(msg slack.Msg) (standupMsg, bool) {
	if msg.Type != "message" || msg.SubType != "" {
		return standupMsg{}, false
	}
	standupText := a.extractMsgRE.ReplaceAllString(msg.Text, "$"+extractMsgGroupName)
	if len(standupText) == len(msg.Text) {
		// Nothing was extracted
		return standupMsg{}, false
	}
	ts, err := parseSlackTimeStamp(msg.Timestamp)
	if err != nil {
		return standupMsg{}, false
	}
	return standupMsg{ts, standupText}, true
}

func (a arriba) retrieveChannelStandup(c conversation) (channelStandup, error) {
	params := slack.NewHistoryParameters()
	params.Count = 1000
	now := time.Now().UTC()
	params.Latest = fmt.Sprintf("%d", now.Unix())
	params.Oldest = fmt.Sprintf("%d", now.AddDate(0, 0, -a.historyDaysLimit).Unix())

	// It would be way more efficient to use slack.SearchMsgs instead
	// of traversing the whole history, but that's not allowed for bots :(
	cstandup := make(channelStandup)
	for {
		logrus.Debugf(
			"Requesting history for conversation %s with parameters %#v",
			c.getID(),
			params)

		history, error := c.getHistory(a.rtm, params)
		if error != nil || history == nil || len(history.Messages) == 0 {
			return cstandup, error
		}

		logrus.Debugf(
			"Got history chunk (from %s to %s, latest %s) for conversation %s",
			history.Messages[len(history.Messages)-1].Msg.Timestamp,
			history.Messages[0].Msg.Timestamp, history.Latest, c.getID())

		for _, msg := range history.Messages {
			if _, ok := cstandup[msg.User]; ok {
				// we already have the latest standup message for this user
				continue
			}
			standupMsg, ok := a.extractChannelStandupMsg(msg.Msg)
			if ok && standupMsg.text != "" {
				cstandup[msg.User] = standupMsg
			}
		}

		if !history.HasMore {
			break
		}
		latestMsg := history.Messages[len(history.Messages)-1]
		params.Latest = latestMsg.Timestamp
		params.Inclusive = false
	}
	return cstandup, nil
}

func (a arriba) retrieveStandups(conversations []conversation) {
	for _, c := range conversations {
		logrus.Infof("Retrieveing standup for conversation #%s (%s)", c.getName(), c.getID())
		cstandup, err := a.retrieveChannelStandup(c)
		if err != nil {
			logrus.Errorf("Can't retrieve channel standup for conversation #%s: %s", c.getName(), err)
		}
		a.standups[c.getID()] = cstandup
		logrus.Infof("Standup for conversation #%s (%s) updated to %#v", c.getName(), c.getID(), cstandup)
	}
}

func (a arriba) getUserName(userID string) string {
	info, err := a.rtm.GetUserInfo(userID)
	userName := "id" + userID
	if err != nil {
		logrus.Errorf("Couldn't get user information for user %s: %s", userID, err)
	} else {
		userName = info.Name
	}
	return userName
}

func (a arriba) removeOldMessages(channelID string) {
	cstandup, ok := a.standups[channelID]
	if !ok {
		return
	}
	oldestAllowed := time.Now().UTC().AddDate(0, 0, -a.historyDaysLimit)
	for userID, msg := range cstandup {
		if msg.ts.Before(oldestAllowed) {
			delete(cstandup, userID)
		}
	}
}

func (a arriba) prettyPrintChannelStandup(cstandup channelStandup) string {
	text := "¡Ándale! ¡Ándale! here's the standup status :tada:\n"
	for _, userID := range cstandup.getKeysByTimestamp() {
		standupMsg := cstandup[userID]
		humanTime := humanize.Time(standupMsg.ts)
		userName := a.getUserName(userID)
		// Inject zero-width unicode character in username to avoid notifying users
		if len(userName) > 1 {
			userName = string(userName[0]) + "\ufeff" + string(userName[1:])
		}
		text += fmt.Sprintf("*%s*: %s _(%s)_\n", userName, standupMsg.text, humanTime)
	}
	return text
}

func (a arriba) sendStatus(channelID string) {
	var statusText string
	if cstandup, ok := a.standups[channelID]; ok && len(cstandup) > 0 {
		statusText = a.prettyPrintChannelStandup(cstandup)
	} else {
		statusText = fmt.Sprintf("No standup messages found\nType a message starting with *@%s* to record your standup message", a.botName)
	}
	a.rtm.SendMessage(a.rtm.NewOutgoingMessage(statusText, channelID))

}

func (a arriba) updateLastStandup(channelID, userID string, msg standupMsg) {
	if _, ok := a.standups[channelID]; !ok {
		a.standups[channelID] = make(channelStandup)
	}
	a.standups[channelID][userID] = msg
	confirmationText := fmt.Sprintf("<@%s>: ¡Yeppa! standup status recorded :taco:", userID)
	a.rtm.SendMessage(a.rtm.NewOutgoingMessage(confirmationText, channelID))
}

func (a *arriba) handleConnectedEvent(ev *slack.ConnectedEvent) {
	if a.botID != "" {
		logrus.Warn("Received unexpected Connected event")
		return
	}
	logrus.Infof(
		"Connected as user %s (%s) to team %s (%s)",
		ev.Info.User.Name,
		ev.Info.User.ID,
		ev.Info.Team.Name,
		ev.Info.Team.ID,
	)
	a.botID = ev.Info.User.ID
	a.botName = ev.Info.User.Name
	a.extractMsgRE = regexp.MustCompile(fmt.Sprintf(extractMsgPattern, a.botID))

	// Retrieve standups for public channels and private groups
	var conversations []conversation
	for _, c := range ev.Info.Channels {
		if c.IsMember {
			conversations = append(conversations, channel(c))
		}
	}
	for _, g := range ev.Info.Groups {
		conversations = append(conversations, group(g))
	}
	a.retrieveStandups(conversations)
}

func (a arriba) handleMessageEvent(ev *slack.MessageEvent) {
	logrus.Debugf("Message received %+v", ev)
	if a.botID == "" {
		logrus.Warn("Received message event before finishing initialization")
		return
	}
	if ev.Channel == "" {
		logrus.Warn("Received message with empty channel")
		return
	}
	switch ev.Channel[0] {
	case 'C', 'G':
		// Public and private (group) channels
		smsg, ok := a.extractChannelStandupMsg(ev.Msg)
		if !ok {
			return
		}
		logrus.Infof("Received standup message in channel %s: %+v", ev.Channel, smsg)
		// Garbage-collect old messages
		a.removeOldMessages(ev.Msg.Channel)
		if smsg.text == "" {
			a.sendStatus(ev.Msg.Channel)
		} else {
			a.updateLastStandup(ev.Msg.Channel, ev.Msg.User, smsg)
		}

	case 'D':
		// Direct messages are not supported yet
	}
}

func (a arriba) run() {
	go a.rtm.ManageConnection()

	for {
		select {
		case msg := <-a.rtm.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.ConnectedEvent:
				a.handleConnectedEvent(ev)
			case *slack.MessageEvent:
				a.handleMessageEvent(ev)
			case *slack.RTMError:
				logrus.Error("Invalid credentials", ev.Error())
			case *slack.InvalidAuthEvent:
				logrus.Error("Invalid credentials")
				os.Exit(1)
			}
		}
	}
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s [ flags ] <SlackAPItoken>\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "You can obtain <SlackAPItoken> from https://<yourteam>.slack.com/services/new/bot\n")
}

func main() {
	var (
		debug            bool
		historyDaysLimit int
	)

	flag.Usage = usage
	flag.BoolVar(&debug, "debug", false, "Print debug information")
	flag.IntVar(&historyDaysLimit, "history-limit", 7, "History limit (in days)")
	flag.Parse()
	if len(flag.Args()) < 1 || historyDaysLimit < 1 {
		usage()
		os.Exit(1)
	}

	logrus.SetOutput(os.Stderr)
	if debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	api := slack.New(flag.Arg(0))
	api.SetDebug(debug)

	newArriba(api.NewRTM(), historyDaysLimit).run()
}
<commit_msg>Reorder imports properly<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"regexp"
	"sort"
	"time"

	humanize "github.com/dustin/go-humanize"
	"github.com/nlopes/slack"
	"github.com/sirupsen/logrus"
)

const extractMsgGroupName = "msg"

// extract suffixes of Slack messages starting with @bot-name
var extractMsgPattern = `(?m)^\s*<@%s>:?\s*(?P<` + extractMsgGroupName + `>.*)$`

// A standupMsg is a Slack message directed to the arriba bot (i.e. with a @botname prefix)
type standupMsg struct {
	ts   time.Time
	text string
}

// channelStandup contains the latest standup message of each user in a Slack channel.
type channelStandup map[string]standupMsg

// sortableChannelStandup is a channelStandup, sortable by the timestamp of its messages
// sortableChannelStandup implements sort.Interface to sort the keys of the channelStandup
type sortableChannelStandup struct {
	keys []string
	cs   channelStandup
}

func (s sortableChannelStandup) Swap(i, j int) { s.keys[i], s.keys[j] = s.keys[j], s.keys[i] }
func (s sortableChannelStandup) Len() int      { return len(s.keys) }
func (s sortableChannelStandup) Less(i, j int) bool {
	return s.cs[s.keys[i]].ts.After(s.cs[s.keys[j]].ts)
}

// getKeysByTimestamp returns the userIDs of the standup ordered by their message timestamp (newer first).
func (cs channelStandup) getKeysByTimestamp() []string {
	keys := make([]string, 0, len(cs))
	for k := range cs {
		keys = append(keys, k)
	}
	scs := sortableChannelStandup{
		cs:   cs,
		keys: keys,
	}
	sort.Sort(scs)
	return scs.keys
}

// standups contains the channelStandup of all Slack channels known to the bot.
type standups map[string]channelStandup

// conversation is a generic way to access the IDs, Names and history of both
// slack.Channel and slack.Group. Unfortunately nlopes/slack doesn't expose the
// underlying common type (groupConversation) and we cannot define methods for
// non-local types, which would allow to make things much cleaner ...
type conversation interface {
	getID() string
	getName() string
	getHistory(*slack.RTM, slack.HistoryParameters) (*slack.History, error)
}

type channel slack.Channel

func (c channel) getID() string   { return c.ID }
func (c channel) getName() string { return c.Name }
func (c channel) getHistory(rtm *slack.RTM, params slack.HistoryParameters) (*slack.History, error) {
	return rtm.GetChannelHistory(c.getID(), params)
}

type group slack.Group

func (g group) getID() string   { return g.ID }
func (g group) getName() string { return g.Name }
func (g group) getHistory(rtm *slack.RTM, params slack.HistoryParameters) (*slack.History, error) {
	return rtm.GetGroupHistory(g.getID(), params)
}

type arriba struct {
	rtm              *slack.RTM
	botID            string
	botName          string
	extractMsgRE     *regexp.Regexp
	historyDaysLimit int
	standups         standups
}

func newArriba(rtm *slack.RTM, historyDaysLimit int) arriba {
	return arriba{
		rtm:              rtm,
		historyDaysLimit: historyDaysLimit,
		standups:         make(standups),
	}
}

func parseSlackTimeStamp(ts string) (time.Time, error) {
	var seconds, milliseconds int64
	_, err := fmt.Sscanf(ts, "%d.%d", &seconds, &milliseconds)
	if err != nil {
		logrus.Warn("Can't parse timestamp ", ts)
		return time.Now(), err
	}
	return time.Unix(seconds, milliseconds*1000), nil
}

// extractStandupMsg parses Slack messages starting with @bot-name
func (a arriba) extractChannelStandupMsg(msg slack.Msg) (standupMsg, bool) {
	if msg.Type != "message" || msg.SubType != "" {
		return standupMsg{}, false
	}
	standupText := a.extractMsgRE.ReplaceAllString(msg.Text, "$"+extractMsgGroupName)
	if len(standupText) == len(msg.Text) {
		// Nothing was extracted
		return standupMsg{}, false
	}
	ts, err := parseSlackTimeStamp(msg.Timestamp)
	if err != nil {
		return standupMsg{}, false
	}
	return standupMsg{ts, standupText}, true
}

func (a arriba) retrieveChannelStandup(c conversation) (channelStandup, error) {
	params := slack.NewHistoryParameters()
	params.Count = 1000
	now := time.Now().UTC()
	params.Latest = fmt.Sprintf("%d", now.Unix())
	params.Oldest = fmt.Sprintf("%d", now.AddDate(0, 0, -a.historyDaysLimit).Unix())

	// It would be way more efficient to use slack.SearchMsgs instead
	// of traversing the whole history, but that's not allowed for bots :(
	cstandup := make(channelStandup)
	for {
		logrus.Debugf(
			"Requesting history for conversation %s with parameters %#v",
			c.getID(),
			params)

		history, error := c.getHistory(a.rtm, params)
		if error != nil || history == nil || len(history.Messages) == 0 {
			return cstandup, error
		}

		logrus.Debugf(
			"Got history chunk (from %s to %s, latest %s) for conversation %s",
			history.Messages[len(history.Messages)-1].Msg.Timestamp,
			history.Messages[0].Msg.Timestamp, history.Latest, c.getID())

		for _, msg := range history.Messages {
			if _, ok := cstandup[msg.User]; ok {
				// we already have the latest standup message for this user
				continue
			}
			standupMsg, ok := a.extractChannelStandupMsg(msg.Msg)
			if ok && standupMsg.text != "" {
				cstandup[msg.User] = standupMsg
			}
		}

		if !history.HasMore {
			break
		}
		latestMsg := history.Messages[len(history.Messages)-1]
		params.Latest = latestMsg.Timestamp
		params.Inclusive = false
	}
	return cstandup, nil
}

func (a arriba) retrieveStandups(conversations []conversation) {
	for _, c := range conversations {
		logrus.Infof("Retrieveing standup for conversation #%s (%s)", c.getName(), c.getID())
		cstandup, err := a.retrieveChannelStandup(c)
		if err != nil {
			logrus.Errorf("Can't retrieve channel standup for conversation #%s: %s", c.getName(), err)
		}
		a.standups[c.getID()] = cstandup
		logrus.Infof("Standup for conversation #%s (%s) updated to %#v", c.getName(), c.getID(), cstandup)
	}
}

func (a arriba) getUserName(userID string) string {
	info, err := a.rtm.GetUserInfo(userID)
	userName := "id" + userID
	if err != nil {
		logrus.Errorf("Couldn't get user information for user %s: %s", userID, err)
	} else {
		userName = info.Name
	}
	return userName
}

func (a arriba) removeOldMessages(channelID string) {
	cstandup, ok := a.standups[channelID]
	if !ok {
		return
	}
	oldestAllowed := time.Now().UTC().AddDate(0, 0, -a.historyDaysLimit)
	for userID, msg := range cstandup {
		if msg.ts.Before(oldestAllowed) {
			delete(cstandup, userID)
		}
	}
}

func (a arriba) prettyPrintChannelStandup(cstandup channelStandup) string {
	text := "¡Ándale! ¡Ándale! here's the standup status :tada:\n"
	for _, userID := range cstandup.getKeysByTimestamp() {
		standupMsg := cstandup[userID]
		humanTime := humanize.Time(standupMsg.ts)
		userName := a.getUserName(userID)
		// Inject zero-width unicode character in username to avoid notifying users
		if len(userName) > 1 {
			userName = string(userName[0]) + "\ufeff" + string(userName[1:])
		}
		text += fmt.Sprintf("*%s*: %s _(%s)_\n", userName, standupMsg.text, humanTime)
	}
	return text
}

func (a arriba) sendStatus(channelID string) {
	var statusText string
	if cstandup, ok := a.standups[channelID]; ok && len(cstandup) > 0 {
		statusText = a.prettyPrintChannelStandup(cstandup)
	} else {
		statusText = fmt.Sprintf("No standup messages found\nType a message starting with *@%s* to record your standup message", a.botName)
	}
	a.rtm.SendMessage(a.rtm.NewOutgoingMessage(statusText, channelID))

}

func (a arriba) updateLastStandup(channelID, userID string, msg standupMsg) {
	if _, ok := a.standups[channelID]; !ok {
		a.standups[channelID] = make(channelStandup)
	}
	a.standups[channelID][userID] = msg
	confirmationText := fmt.Sprintf("<@%s>: ¡Yeppa! standup status recorded :taco:", userID)
	a.rtm.SendMessage(a.rtm.NewOutgoingMessage(confirmationText, channelID))
}

func (a *arriba) handleConnectedEvent(ev *slack.ConnectedEvent) {
	if a.botID != "" {
		logrus.Warn("Received unexpected Connected event")
		return
	}
	logrus.Infof(
		"Connected as user %s (%s) to team %s (%s)",
		ev.Info.User.Name,
		ev.Info.User.ID,
		ev.Info.Team.Name,
		ev.Info.Team.ID,
	)
	a.botID = ev.Info.User.ID
	a.botName = ev.Info.User.Name
	a.extractMsgRE = regexp.MustCompile(fmt.Sprintf(extractMsgPattern, a.botID))

	// Retrieve standups for public channels and private groups
	var conversations []conversation
	for _, c := range ev.Info.Channels {
		if c.IsMember {
			conversations = append(conversations, channel(c))
		}
	}
	for _, g := range ev.Info.Groups {
		conversations = append(conversations, group(g))
	}
	a.retrieveStandups(conversations)
}

func (a arriba) handleMessageEvent(ev *slack.MessageEvent) {
	logrus.Debugf("Message received %+v", ev)
	if a.botID == "" {
		logrus.Warn("Received message event before finishing initialization")
		return
	}
	if ev.Channel == "" {
		logrus.Warn("Received message with empty channel")
		return
	}
	switch ev.Channel[0] {
	case 'C', 'G':
		// Public and private (group) channels
		smsg, ok := a.extractChannelStandupMsg(ev.Msg)
		if !ok {
			return
		}
		logrus.Infof("Received standup message in channel %s: %+v", ev.Channel, smsg)
		// Garbage-collect old messages
		a.removeOldMessages(ev.Msg.Channel)
		if smsg.text == "" {
			a.sendStatus(ev.Msg.Channel)
		} else {
			a.updateLastStandup(ev.Msg.Channel, ev.Msg.User, smsg)
		}

	case 'D':
		// Direct messages are not supported yet
	}
}

func (a arriba) run() {
	go a.rtm.ManageConnection()

	for {
		select {
		case msg := <-a.rtm.IncomingEvents:
			switch ev := msg.Data.(type) {
			case *slack.ConnectedEvent:
				a.handleConnectedEvent(ev)
			case *slack.MessageEvent:
				a.handleMessageEvent(ev)
			case *slack.RTMError:
				logrus.Error("Invalid credentials", ev.Error())
			case *slack.InvalidAuthEvent:
				logrus.Error("Invalid credentials")
				os.Exit(1)
			}
		}
	}
}

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s [ flags ] <SlackAPItoken>\n", os.Args[0])
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "You can obtain <SlackAPItoken> from https://<yourteam>.slack.com/services/new/bot\n")
}

func main() {
	var (
		debug            bool
		historyDaysLimit int
	)

	flag.Usage = usage
	flag.BoolVar(&debug, "debug", false, "Print debug information")
	flag.IntVar(&historyDaysLimit, "history-limit", 7, "History limit (in days)")
	flag.Parse()
	if len(flag.Args()) < 1 || historyDaysLimit < 1 {
		usage()
		os.Exit(1)
	}

	logrus.SetOutput(os.Stderr)
	if debug {
		logrus.SetLevel(logrus.DebugLevel)
	}

	api := slack.New(flag.Arg(0))
	api.SetDebug(debug)

	newArriba(api.NewRTM(), historyDaysLimit).run()
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/codahale/hdrhistogram"
)

type MeasuredResponse struct {
	sz      uint64
	code    int
	latency int64
	timeout bool
	err     error
}

func newClient(
	compress bool,
	https bool,
	reuse bool,
	maxConn uint,
) *http.Client {
	tr := http.Transport{
		DisableCompression:  !compress,
		DisableKeepAlives:   !reuse,
		MaxIdleConnsPerHost: int(maxConn),
	}
	if https {
		tr.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	}
	return &http.Client{Transport: &tr}
}

func sendRequest(
	client *http.Client,
	url *url.URL,
	host *string,
	reqID uint64,
	received chan *MeasuredResponse,
) {
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		fmt.Fprintln(os.Stderr, "\n")
	}
	req.Host = *host
	req.Header.Add("Sc-Req-Id", fmt.Sprintf("%d", reqID))

	// FIX: find a way to measure latency with the http client.
	start := time.Now()
	response, err := client.Do(req)

	elapsed := time.Since(start)
	if err != nil {
		received <- &MeasuredResponse{0, 0, 0, false, err}
	} else {
		sz, _ := io.Copy(ioutil.Discard, response.Body)
		response.Body.Close()
		received <- &MeasuredResponse{uint64(sz), response.StatusCode, elapsed.Nanoseconds(), false, nil}
	}
}

func exUsage(msg string) {
	fmt.Fprintf(os.Stderr, "%s\n", msg)
	fmt.Fprintf(os.Stderr, "Usage: %s [flags]\n", path.Base(os.Args[0]))
	flag.PrintDefaults()
	os.Exit(64)
}

// To achieve a target qps, we need to wait this many Nanoseconds
// between actions.
func CalcTimeToWait(qps *int) time.Duration {
	return time.Duration(int(time.Second) / *qps)
}

var reqID = uint64(0)

var shouldFinish = false
var shouldFinishLock sync.RWMutex

// Signals the system to stop sending traffic and clean up after itself.
func finishSendingTraffic() {
	shouldFinishLock.Lock()
	shouldFinish = true
	shouldFinishLock.Unlock()
}

func main() {
	qps := flag.Int("qps", 1, "QPS to send to backends per request thread")
	concurrency := flag.Uint("concurrency", 1, "Number of request threads")
	host := flag.String("host", "web", "value of Host header to set")
	urldest := flag.String("url", "http://localhost:4140/", "Destination url")
	interval := flag.Duration("interval", 10*time.Second, "reporting interval")
	reuse := flag.Bool("reuse", false, "reuse connections")
	compress := flag.Bool("compress", false, "use compression")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags]\n", path.Base(os.Args[0]))
		flag.PrintDefaults()
	}

	flag.Parse()

	if *qps < 1 {
		exUsage("qps must be at least 1")
	}

	if *concurrency < 1 {
		exUsage("concurrency must be at least 1")
	}

	hosts := strings.Split(*host, ",")

	dstURL, err := url.Parse(*urldest)
	if err != nil {
		exUsage(fmt.Sprintf("invalid URL: '%s': %s\n", urldest, err.Error()))
	}

	// Repsonse tracking metadata.
	count := uint64(0)
	size := uint64(0)
	good := uint64(0)
	bad := uint64(0)
	failed := uint64(0)
	// from 0 to 1 minute in nanoseconds
	// FIX: verify that these buckets work correctly for our use case.
	hist := hdrhistogram.New(0, 60000000000, 5)
	received := make(chan *MeasuredResponse)
	timeout := time.After(*interval)
	timeToWait := CalcTimeToWait(qps)

	doTLS := dstURL.Scheme == "https"
	client := newClient(*compress, doTLS, *reuse, *concurrency)
	var sendTraffic sync.WaitGroup

	for i := uint(0); i < *concurrency; i++ {
		ticker := time.NewTicker(timeToWait)
		go func() {
			sendTraffic.Add(1)
			for _ = range ticker.C {
				shouldFinishLock.RLock()
				if !shouldFinish {
					shouldFinishLock.RUnlock()
					sendRequest(client, dstURL, &hosts[rand.Intn(len(hosts))], atomic.AddUint64(&reqID, 1), received)
				} else {
					shouldFinishLock.RUnlock()
					sendTraffic.Done()
					return
				}
			}
		}()
	}

	cleanup := make(chan os.Signal)
	signal.Notify(cleanup, syscall.SIGINT)

	for {
		select {
		case <-cleanup:
			finishSendingTraffic()
			go func() {
				// Don't Wait() in the event loop or else we'll block the workers
				// from draining.
				sendTraffic.Wait()
				os.Exit(1)
			}()
		case t := <-timeout:
			// Periodically print stats about the request load.
			fmt.Printf("%s %6d/%1d/%1d requests %6d kilobytes %s [%3d %3d %3d %4d ]\n",
				t.Format(time.RFC3339),
				good,
				bad,
				failed,
				(size / 1024),
				interval,
				hist.ValueAtQuantile(50)/1000000,
				hist.ValueAtQuantile(95)/1000000,
				hist.ValueAtQuantile(99)/1000000,
				hist.ValueAtQuantile(999)/1000000)
			count = 0
			size = 0
			good = 0
			bad = 0
			failed = 0
			hist.Reset()
			timeout = time.After(*interval)
		case managedResp := <-received:
			count++
			if managedResp.err != nil {
				fmt.Fprintln(os.Stderr, managedResp.err)
				failed++
			} else {
				size += managedResp.sz
				if managedResp.code >= 200 && managedResp.code < 500 {
					good++
				} else {
					bad++
				}
				hist.RecordValue(managedResp.latency)
			}
		}
	}
}
<commit_msg>fixup<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"path"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/codahale/hdrhistogram"
)

type MeasuredResponse struct {
	sz      uint64
	code    int
	latency int64
	timeout bool
	err     error
}

func newClient(
	compress bool,
	https bool,
	reuse bool,
	maxConn uint,
) *http.Client {
	tr := http.Transport{
		DisableCompression:  !compress,
		DisableKeepAlives:   !reuse,
		MaxIdleConnsPerHost: int(maxConn),
	}
	if https {
		tr.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	}
	return &http.Client{Transport: &tr}
}

func sendRequest(
	client *http.Client,
	url *url.URL,
	host *string,
	reqID uint64,
	received chan *MeasuredResponse,
) {
	req, err := http.NewRequest("GET", url.String(), nil)
	if err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		fmt.Fprintln(os.Stderr, "\n")
	}
	req.Host = *host
	req.Header.Add("Sc-Req-Id", strconv.FormatUint(reqID, 10))

	// FIX: find a way to measure latency with the http client.
	start := time.Now()
	response, err := client.Do(req)

	elapsed := time.Since(start)
	if err != nil {
		received <- &MeasuredResponse{0, 0, 0, false, err}
	} else {
		sz, _ := io.Copy(ioutil.Discard, response.Body)
		response.Body.Close()
		received <- &MeasuredResponse{uint64(sz), response.StatusCode, elapsed.Nanoseconds(), false, nil}
	}
}

func exUsage(msg string) {
	fmt.Fprintf(os.Stderr, "%s\n", msg)
	fmt.Fprintf(os.Stderr, "Usage: %s [flags]\n", path.Base(os.Args[0]))
	flag.PrintDefaults()
	os.Exit(64)
}

// To achieve a target qps, we need to wait this many Nanoseconds
// between actions.
func CalcTimeToWait(qps *int) time.Duration {
	return time.Duration(int(time.Second) / *qps)
}

var reqID = uint64(0)

var shouldFinish = false
var shouldFinishLock sync.RWMutex

// Signals the system to stop sending traffic and clean up after itself.
func finishSendingTraffic() {
	shouldFinishLock.Lock()
	shouldFinish = true
	shouldFinishLock.Unlock()
}

func main() {
	qps := flag.Int("qps", 1, "QPS to send to backends per request thread")
	concurrency := flag.Uint("concurrency", 1, "Number of request threads")
	host := flag.String("host", "web", "value of Host header to set")
	urldest := flag.String("url", "http://localhost:4140/", "Destination url")
	interval := flag.Duration("interval", 10*time.Second, "reporting interval")
	reuse := flag.Bool("reuse", false, "reuse connections")
	compress := flag.Bool("compress", false, "use compression")

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags]\n", path.Base(os.Args[0]))
		flag.PrintDefaults()
	}

	flag.Parse()

	if *qps < 1 {
		exUsage("qps must be at least 1")
	}

	if *concurrency < 1 {
		exUsage("concurrency must be at least 1")
	}

	hosts := strings.Split(*host, ",")

	dstURL, err := url.Parse(*urldest)
	if err != nil {
		exUsage(fmt.Sprintf("invalid URL: '%s': %s\n", urldest, err.Error()))
	}

	// Repsonse tracking metadata.
	count := uint64(0)
	size := uint64(0)
	good := uint64(0)
	bad := uint64(0)
	failed := uint64(0)
	// from 0 to 1 minute in nanoseconds
	// FIX: verify that these buckets work correctly for our use case.
	hist := hdrhistogram.New(0, 60000000000, 5)
	received := make(chan *MeasuredResponse)
	timeout := time.After(*interval)
	timeToWait := CalcTimeToWait(qps)

	doTLS := dstURL.Scheme == "https"
	client := newClient(*compress, doTLS, *reuse, *concurrency)
	var sendTraffic sync.WaitGroup

	for i := uint(0); i < *concurrency; i++ {
		ticker := time.NewTicker(timeToWait)
		go func() {
			sendTraffic.Add(1)
			for _ = range ticker.C {
				shouldFinishLock.RLock()
				if !shouldFinish {
					shouldFinishLock.RUnlock()
					sendRequest(client, dstURL, &hosts[rand.Intn(len(hosts))], atomic.AddUint64(&reqID, 1), received)
				} else {
					shouldFinishLock.RUnlock()
					sendTraffic.Done()
					return
				}
			}
		}()
	}

	cleanup := make(chan os.Signal)
	signal.Notify(cleanup, syscall.SIGINT)

	for {
		select {
		case <-cleanup:
			finishSendingTraffic()
			go func() {
				// Don't Wait() in the event loop or else we'll block the workers
				// from draining.
				sendTraffic.Wait()
				os.Exit(1)
			}()
		case t := <-timeout:
			// Periodically print stats about the request load.
			fmt.Printf("%s %6d/%1d/%1d requests %6d kilobytes %s [%3d %3d %3d %4d ]\n",
				t.Format(time.RFC3339),
				good,
				bad,
				failed,
				(size / 1024),
				interval,
				hist.ValueAtQuantile(50)/1000000,
				hist.ValueAtQuantile(95)/1000000,
				hist.ValueAtQuantile(99)/1000000,
				hist.ValueAtQuantile(999)/1000000)
			count = 0
			size = 0
			good = 0
			bad = 0
			failed = 0
			hist.Reset()
			timeout = time.After(*interval)
		case managedResp := <-received:
			count++
			if managedResp.err != nil {
				fmt.Fprintln(os.Stderr, managedResp.err)
				failed++
			} else {
				size += managedResp.sz
				if managedResp.code >= 200 && managedResp.code < 500 {
					good++
				} else {
					bad++
				}
				hist.RecordValue(managedResp.latency)
			}
		}
	}
}
<|endoftext|>