<commit_before>package model_loader

import (
	"encoding/json"
	"fmt"
	"github.com/centurylinkcloud/clc-go-cli/parser"
	"reflect"
	"time"
)

const (
	timeFormat = "2006-01-02 15:04:05"
)

func LoadModel(parsedArgs map[string]interface{}, inputModel interface{}) error {
	metaModel := reflect.ValueOf(inputModel)
	if metaModel.Kind() != reflect.Ptr {
		return fmt.Errorf("Input model must be passed by pointer.")
	}
	for k, v := range parsedArgs {
		field, err := getFieldByName(metaModel, k)
		if err != nil {
			return err
		}
		err = loadValue(k, v, field)
		if err != nil {
			return err
		}
	}
	return nil
}

func loadValue(key string, arg interface{}, field reflect.Value) error {
	switch field.Interface().(type) {
	case int64:
		if argInt, isInt := arg.(int64); !isInt {
			return fmt.Errorf("Type mismatch: %s value must be integer.", key)
		} else {
			field.SetInt(argInt)
			return nil
		}
	case float64:
		if argFloat, isFloat := arg.(float64); !isFloat {
			return fmt.Errorf("Type mismatch: %s value must be float.", key)
		} else {
			field.SetFloat(argFloat)
			return nil
		}
	case time.Time:
		if argTime, err := time.Parse(timeFormat, arg.(string)); err != nil {
			return fmt.Errorf("Type mismatch: %s value must be datetime in `YYYY-MM-DD hh:mm:ss` format.", key)
		} else {
			field.Set(reflect.ValueOf(argTime))
			return nil
		}
	case bool:
		if arg == "true" {
			field.SetBool(true)
		} else if arg == "false" {
			field.SetBool(false)
		} else {
			return fmt.Errorf("Type mismatch: %s value must be either true or false.", key)
		}
		return nil
	case string:
		field.SetString(arg.(string))
		return nil
	}
	if isStruct(field) {
		argStruct, err := parseStruct(arg)
		if err != nil {
			return err
		}
		for k, v := range argStruct {
			nestedField, err := getFieldByName(field.Addr(), k)
			if err != nil {
				return err
			}
			err = loadValue(k, v, nestedField)
			if err != nil {
				return err
			}
		}
		return nil
	} else if isSlice(field) {
		argSlice, err := parseSlice(arg)
		if err != nil {
			return err
		}
		for _, v := range argSlice {
			elementPtr := getEmplySliceType(field)
			err = loadValue(key, v, elementPtr.Elem())
			if err != nil {
				return err
			}
			field.Set(reflect.Append(field, elementPtr.Elem()))
		}
		return nil
	}
	return fmt.Errorf("Unsupported field type %s", field.Kind())
}

func getFieldByName(model reflect.Value, name string) (reflect.Value, error) {
	field := model.Elem().FieldByName(name)
	if !field.IsValid() {
		return reflect.ValueOf(nil), fmt.Errorf("Field `%s` does not exist.", name)
	}
	return field, nil
}

func isStruct(model reflect.Value) bool {
	return model.Kind() == reflect.Struct
}

func isSlice(model reflect.Value) bool {
	return model.Kind() == reflect.Slice
}

// Parses an object of type map[string]interface{} either from JSON or from a=b,c=d,.. notation.
// Also, calls NormalizeKeys with the parsed object.
// If arg is already of type map[string]interface{} returns it as is.
func parseStruct(arg interface{}) (map[string]interface{}, error) {
	if argMap, isMap := arg.(map[string]interface{}); isMap {
		return argMap, nil
	}
	parsed := make(map[string]interface{}, 0)
	if err := json.Unmarshal([]byte(arg.(string)), &parsed); err == nil {
		parser.NormalizeKeys(parsed)
		return parsed, nil
	}
	// TODO parse a=b,c=d,.. notation
	return nil, fmt.Errorf("`%s` is neither in JSON nor in key=value,.. format.", arg.(string))
}

// Parses an object of type []interface{} either from JSON.
// Also, calls NormalizeKeys with the parsed object.
// If arg is already of type []interface{} returns it as is.
func parseSlice(arg interface{}) ([]interface{}, error) {
	if argSlice, isSlice := arg.([]interface{}); isSlice {
		return argSlice, nil
	}
	parsed := make([]interface{}, 0)
	if err := json.Unmarshal([]byte(arg.(string)), &parsed); err == nil {
		parser.NormalizeKeys(parsed)
		return parsed, nil
	}
	return nil, fmt.Errorf("`%s` is neither in JSON nor in key=value,.. format.", arg.(string))
}

func getEmplySliceType(slice reflect.Value) reflect.Value {
	return reflect.New(slice.Type().Elem())
}
<commit_msg>add the a=b,c=d,.. notation parsing step; all existing model loader tests pass<commit_after>package model_loader

import (
	"encoding/json"
	"fmt"
	"github.com/centurylinkcloud/clc-go-cli/parser"
	"reflect"
	"time"
)

const (
	timeFormat = "2006-01-02 15:04:05"
)

func LoadModel(parsedArgs map[string]interface{}, inputModel interface{}) error {
	metaModel := reflect.ValueOf(inputModel)
	if metaModel.Kind() != reflect.Ptr {
		return fmt.Errorf("Input model must be passed by pointer.")
	}
	for k, v := range parsedArgs {
		field, err := getFieldByName(metaModel, k)
		if err != nil {
			return err
		}
		err = loadValue(k, v, field)
		if err != nil {
			return err
		}
	}
	return nil
}

func loadValue(key string, arg interface{}, field reflect.Value) error {
	switch field.Interface().(type) {
	case int64:
		if argInt, isInt := arg.(int64); !isInt {
			return fmt.Errorf("Type mismatch: %s value must be integer.", key)
		} else {
			field.SetInt(argInt)
			return nil
		}
	case float64:
		if argFloat, isFloat := arg.(float64); !isFloat {
			return fmt.Errorf("Type mismatch: %s value must be float.", key)
		} else {
			field.SetFloat(argFloat)
			return nil
		}
	case time.Time:
		if argTime, err := time.Parse(timeFormat, arg.(string)); err != nil {
			return fmt.Errorf("Type mismatch: %s value must be datetime in `YYYY-MM-DD hh:mm:ss` format.", key)
		} else {
			field.Set(reflect.ValueOf(argTime))
			return nil
		}
	case bool:
		if arg == "true" {
			field.SetBool(true)
		} else if arg == "false" {
			field.SetBool(false)
		} else {
			return fmt.Errorf("Type mismatch: %s value must be either true or false.", key)
		}
		return nil
	case string:
		field.SetString(arg.(string))
		return nil
	}
	if isStruct(field) {
		argStruct, err := parseStruct(arg)
		if err != nil {
			return err
		}
		for k, v := range argStruct {
			nestedField, err := getFieldByName(field.Addr(), k)
			if err != nil {
				return err
			}
			err = loadValue(k, v, nestedField)
			if err != nil {
				return err
			}
		}
		return nil
	} else if isSlice(field) {
		argSlice, err := parseSlice(arg)
		if err != nil {
			return err
		}
		for _, v := range argSlice {
			elementPtr := getEmplySliceType(field)
			err = loadValue(key, v, elementPtr.Elem())
			if err != nil {
				return err
			}
			field.Set(reflect.Append(field, elementPtr.Elem()))
		}
		return nil
	}
	return fmt.Errorf("Unsupported field type %s", field.Kind())
}

func getFieldByName(model reflect.Value, name string) (reflect.Value, error) {
	field := model.Elem().FieldByName(name)
	if !field.IsValid() {
		return reflect.ValueOf(nil), fmt.Errorf("Field `%s` does not exist.", name)
	}
	return field, nil
}

func isStruct(model reflect.Value) bool {
	return model.Kind() == reflect.Struct
}

func isSlice(model reflect.Value) bool {
	return model.Kind() == reflect.Slice
}

// Parses an object of type map[string]interface{} either from JSON or from a=b,c=d,.. notation.
// Also, calls NormalizeKeys with the parsed object.
// If arg is already of type map[string]interface{} returns it as is.
func parseStruct(arg interface{}) (map[string]interface{}, error) {
	if argMap, isMap := arg.(map[string]interface{}); isMap {
		return argMap, nil
	}
	parsed := make(map[string]interface{}, 0)
	if err := json.Unmarshal([]byte(arg.(string)), &parsed); err == nil {
		parser.NormalizeKeys(parsed)
		return parsed, nil
	}
	if parsed, err := parser.ParseObject(arg.(string)); err == nil {
		return parsed, nil
	}
	return nil, fmt.Errorf("`%s` is neither in JSON nor in key=value,.. format.", arg.(string))
}

// Parses an object of type []interface{} either from JSON.
// Also, calls NormalizeKeys with the parsed object.
// If arg is already of type []interface{} returns it as is.
func parseSlice(arg interface{}) ([]interface{}, error) {
	if argSlice, isSlice := arg.([]interface{}); isSlice {
		return argSlice, nil
	}
	parsed := make([]interface{}, 0)
	if err := json.Unmarshal([]byte(arg.(string)), &parsed); err == nil {
		parser.NormalizeKeys(parsed)
		return parsed, nil
	}
	return nil, fmt.Errorf("`%s` is neither in JSON nor in key=value,.. format.", arg.(string))
}

func getEmplySliceType(slice reflect.Value) reflect.Value {
	return reflect.New(slice.Type().Elem())
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"go/ast"
)

var osopenFix = fix{
	"osopen",
	osopen,
	`Adapt os.Open calls to new, easier API and rename O_CREAT O_CREATE.

	http://codereview.appspot.com/4357052
`,
}

func init() {
	register(osopenFix)
}

func osopen(f *ast.File) bool {
	if !imports(f, "os") {
		return false
	}

	fixed := false
	walk(f, func(n interface{}) {
		// Rename O_CREAT to O_CREATE.
		if expr, ok := n.(ast.Expr); ok && isPkgDot(expr, "os", "O_CREAT") {
			expr.(*ast.SelectorExpr).Sel.Name = "O_CREATE"
			return
		}

		// Fix up calls to Open.
		call, ok := n.(*ast.CallExpr)
		if !ok || len(call.Args) != 3 {
			return
		}
		if !isPkgDot(call.Fun, "os", "Open") {
			return
		}
		sel := call.Fun.(*ast.SelectorExpr)
		args := call.Args
		// os.Open(a, os.O_RDONLY, c) -> os.Open(a)
		if isPkgDot(args[1], "os", "O_RDONLY") || isPkgDot(args[1], "syscall", "O_RDONLY") {
			call.Args = call.Args[0:1]
			fixed = true
			return
		}
		// os.Open(a, createlike_flags, c) -> os.Create(a, c)
		if isCreateFlag(args[1]) {
			sel.Sel.Name = "Create"
			if !isSimplePerm(args[2]) {
				warn(sel.Pos(), "rewrote os.Open to os.Create with permission not 0666")
			}
			call.Args = args[0:1]
			fixed = true
			return
		}
		// Fallback: os.Open(a, b, c) -> os.OpenFile(a, b, c)
		sel.Sel.Name = "OpenFile"
		fixed = true
	})
	return fixed
}

func isCreateFlag(flag ast.Expr) bool {
	foundCreate := false
	foundTrunc := false
	// OR'ing of flags: is O_CREATE on?  + or | would be fine; we just look for os.O_CREATE
	// and don't worry about the actual opeator.
	p := flag.Pos()
	for {
		lhs := flag
		expr, isBinary := flag.(*ast.BinaryExpr)
		if isBinary {
			lhs = expr.Y
		}
		sel, ok := lhs.(*ast.SelectorExpr)
		if !ok || !isTopName(sel.X, "os") {
			return false
		}
		switch sel.Sel.Name {
		case "O_CREATE":
			foundCreate = true
		case "O_TRUNC":
			foundTrunc = true
		case "O_RDONLY", "O_WRONLY", "O_RDWR":
			// okay
		default:
			// Unexpected flag, like O_APPEND or O_EXCL.
			// Be conservative and do not rewrite.
			return false
		}
		if !isBinary {
			break
		}
		flag = expr.X
	}
	if !foundCreate {
		return false
	}
	if !foundTrunc {
		warn(p, "rewrote os.Open with O_CREATE but not O_TRUNC to os.Create")
	}
	return foundCreate
}

func isSimplePerm(perm ast.Expr) bool {
	basicLit, ok := perm.(*ast.BasicLit)
	if !ok {
		return false
	}
	switch basicLit.Value {
	case "0666":
		return true
	}
	return false
}
<commit_msg>gofix: fix embarrassing typo in osopen.go<commit_after>// Copyright 2011 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"go/ast"
)

var osopenFix = fix{
	"osopen",
	osopen,
	`Adapt os.Open calls to new, easier API and rename O_CREAT O_CREATE.

	http://codereview.appspot.com/4357052
`,
}

func init() {
	register(osopenFix)
}

func osopen(f *ast.File) bool {
	if !imports(f, "os") {
		return false
	}

	fixed := false
	walk(f, func(n interface{}) {
		// Rename O_CREAT to O_CREATE.
		if expr, ok := n.(ast.Expr); ok && isPkgDot(expr, "os", "O_CREAT") {
			expr.(*ast.SelectorExpr).Sel.Name = "O_CREATE"
			return
		}

		// Fix up calls to Open.
		call, ok := n.(*ast.CallExpr)
		if !ok || len(call.Args) != 3 {
			return
		}
		if !isPkgDot(call.Fun, "os", "Open") {
			return
		}
		sel := call.Fun.(*ast.SelectorExpr)
		args := call.Args
		// os.Open(a, os.O_RDONLY, c) -> os.Open(a)
		if isPkgDot(args[1], "os", "O_RDONLY") || isPkgDot(args[1], "syscall", "O_RDONLY") {
			call.Args = call.Args[0:1]
			fixed = true
			return
		}
		// os.Open(a, createlike_flags, c) -> os.Create(a, c)
		if isCreateFlag(args[1]) {
			sel.Sel.Name = "Create"
			if !isSimplePerm(args[2]) {
				warn(sel.Pos(), "rewrote os.Open to os.Create with permission not 0666")
			}
			call.Args = args[0:1]
			fixed = true
			return
		}
		// Fallback: os.Open(a, b, c) -> os.OpenFile(a, b, c)
		sel.Sel.Name = "OpenFile"
		fixed = true
	})
	return fixed
}

func isCreateFlag(flag ast.Expr) bool {
	foundCreate := false
	foundTrunc := false
	// OR'ing of flags: is O_CREATE on?  + or | would be fine; we just look for os.O_CREATE
	// and don't worry about the actual operator.
	p := flag.Pos()
	for {
		lhs := flag
		expr, isBinary := flag.(*ast.BinaryExpr)
		if isBinary {
			lhs = expr.Y
		}
		sel, ok := lhs.(*ast.SelectorExpr)
		if !ok || !isTopName(sel.X, "os") {
			return false
		}
		switch sel.Sel.Name {
		case "O_CREATE":
			foundCreate = true
		case "O_TRUNC":
			foundTrunc = true
		case "O_RDONLY", "O_WRONLY", "O_RDWR":
			// okay
		default:
			// Unexpected flag, like O_APPEND or O_EXCL.
			// Be conservative and do not rewrite.
			return false
		}
		if !isBinary {
			break
		}
		flag = expr.X
	}
	if !foundCreate {
		return false
	}
	if !foundTrunc {
		warn(p, "rewrote os.Open with O_CREATE but not O_TRUNC to os.Create")
	}
	return foundCreate
}

func isSimplePerm(perm ast.Expr) bool {
	basicLit, ok := perm.(*ast.BasicLit)
	if !ok {
		return false
	}
	switch basicLit.Value {
	case "0666":
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"math/big"
	"os"
	"strings"
	"time"

	"ledger"

	"github.com/jbrukh/bayesian"
)

func usage() {
	fmt.Printf("Usage: %s -f <ledger-file> <account> <csv file>\n", os.Args[0])
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	var ledgerFileName string
	var accountSubstring, csvFileName, csvDateFormat string
	var negateAmount bool

	flag.BoolVar(&negateAmount, "neg", false, "Negate amount column value.")
	flag.StringVar(&ledgerFileName, "f", "", "Ledger file name (*Required).")
	flag.StringVar(&csvDateFormat, "date-format", "01/02/2006", "Date format.")
	flag.Parse()

	args := flag.Args()
	if len(args) != 2 {
		usage()
	} else {
		accountSubstring = args[0]
		csvFileName = args[1]
	}

	csvFileReader, err := os.Open(csvFileName)
	if err != nil {
		fmt.Println("CSV: ", err)
		return
	}
	defer csvFileReader.Close()

	ledgerFileReader, err := os.Open(ledgerFileName)
	if err != nil {
		fmt.Println("Ledger: ", err)
		return
	}
	defer ledgerFileReader.Close()

	generalLedger, parseError := ledger.ParseLedger(ledgerFileReader)
	if parseError != nil {
		fmt.Println(parseError)
		return
	}

	var matchingAccount string
	matchingAccounts := ledger.GetBalances(generalLedger, []string{accountSubstring})
	if len(matchingAccounts) < 1 {
		fmt.Println("Unable to find matching account.")
		return
	} else {
		matchingAccount = matchingAccounts[len(matchingAccounts)-1].Name
	}

	allAccounts := ledger.GetBalances(generalLedger, []string{})

	csvReader := csv.NewReader(csvFileReader)
	csvRecords, _ := csvReader.ReadAll()

	classes := make([]bayesian.Class, len(allAccounts))
	for i, bal := range allAccounts {
		classes[i] = bayesian.Class(bal.Name)
	}
	classifier := bayesian.NewClassifier(classes...)
	for _, tran := range generalLedger {
		payeeWords := strings.Split(tran.Payee, " ")
		for _, accChange := range tran.AccountChanges {
			if strings.Contains(accChange.Name, "Expense") {
				classifier.Learn(payeeWords, bayesian.Class(accChange.Name))
			}
		}
	}

	// Find columns from header
	var dateColumn, payeeColumn, amountColumn int
	dateColumn, payeeColumn, amountColumn = -1, -1, -1
	for fieldIndex, fieldName := range csvRecords[0] {
		fieldName = strings.ToLower(fieldName)
		if strings.Contains(fieldName, "date") {
			dateColumn = fieldIndex
		} else if strings.Contains(fieldName, "description") {
			payeeColumn = fieldIndex
		} else if strings.Contains(fieldName, "amount") {
			amountColumn = fieldIndex
		}
	}

	if dateColumn < 0 || payeeColumn < 0 || amountColumn < 0 {
		fmt.Println("Unable to find columns required from header field names.")
		return
	}

	expenseAccount := ledger.Account{Name: "unknown:unknown", Balance: new(big.Rat)}
	csvAccount := ledger.Account{Name: matchingAccount, Balance: new(big.Rat)}
	for _, record := range csvRecords[1:] {
		inputPayeeWords := strings.Split(record[payeeColumn], " ")
		csvDate, _ := time.Parse(csvDateFormat, record[dateColumn])
		if !existingTransaction(generalLedger, csvDate, inputPayeeWords[0]) {
			// Classify into expense account
			_, likely, _ := classifier.LogScores(inputPayeeWords)
			if likely >= 0 {
				expenseAccount.Name = string(classifier.Classes[likely])
			}

			// Negate amount if required
			expenseAccount.Balance.SetString(record[amountColumn])
			if negateAmount {
				expenseAccount.Balance.Neg(expenseAccount.Balance)
			}

			// Csv amount is the negative of the expense amount
			csvAccount.Balance.Neg(expenseAccount.Balance)

			// Create valid transaction for print in ledger format
			trans := &ledger.Transaction{Date: csvDate, Payee: record[payeeColumn]}
			trans.AccountChanges = []ledger.Account{csvAccount, expenseAccount}
			PrintTransaction(trans, 80)
		}
	}
}

func existingTransaction(generalLedger []*ledger.Transaction, transDate time.Time, payee string) bool {
	for _, trans := range generalLedger {
		if trans.Date == transDate && strings.HasPrefix(trans.Payee, payee) {
			return true
		}
	}
	return false
}
<commit_msg>Add extra features for importing other files<commit_after>package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"math/big"
	"os"
	"strings"
	"time"
	"unicode/utf8"

	"ledger"

	"github.com/jbrukh/bayesian"
)

func usage() {
	fmt.Printf("Usage: %s -f <ledger-file> <account> <csv file>\n", os.Args[0])
	flag.PrintDefaults()
	os.Exit(1)
}

func main() {
	var ledgerFileName string
	var accountSubstring, csvFileName, csvDateFormat string
	var negateAmount bool
	var fieldDelimiter string

	flag.BoolVar(&negateAmount, "neg", false, "Negate amount column value.")
	flag.StringVar(&ledgerFileName, "f", "", "Ledger file name (*Required).")
	flag.StringVar(&csvDateFormat, "date-format", "01/02/2006", "Date format.")
	flag.StringVar(&fieldDelimiter, "delimiter", ",", "Field delimiter.")
	flag.Parse()

	args := flag.Args()
	if len(args) != 2 {
		usage()
	} else {
		accountSubstring = args[0]
		csvFileName = args[1]
	}

	csvFileReader, err := os.Open(csvFileName)
	if err != nil {
		fmt.Println("CSV: ", err)
		return
	}
	defer csvFileReader.Close()

	ledgerFileReader, err := os.Open(ledgerFileName)
	if err != nil {
		fmt.Println("Ledger: ", err)
		return
	}
	defer ledgerFileReader.Close()

	generalLedger, parseError := ledger.ParseLedger(ledgerFileReader)
	if parseError != nil {
		fmt.Println(parseError)
		return
	}

	var matchingAccount string
	matchingAccounts := ledger.GetBalances(generalLedger, []string{accountSubstring})
	if len(matchingAccounts) < 1 {
		fmt.Println("Unable to find matching account.")
		return
	} else {
		matchingAccount = matchingAccounts[len(matchingAccounts)-1].Name
	}

	allAccounts := ledger.GetBalances(generalLedger, []string{})

	csvReader := csv.NewReader(csvFileReader)
	csvReader.Comma, _ = utf8.DecodeRuneInString(fieldDelimiter)
	csvRecords, _ := csvReader.ReadAll()

	classes := make([]bayesian.Class, len(allAccounts))
	for i, bal := range allAccounts {
		classes[i] = bayesian.Class(bal.Name)
	}
	classifier := bayesian.NewClassifier(classes...)
	for _, tran := range generalLedger {
		payeeWords := strings.Split(tran.Payee, " ")
		for _, accChange := range tran.AccountChanges {
			if strings.Contains(accChange.Name, "Expense") {
				classifier.Learn(payeeWords, bayesian.Class(accChange.Name))
			}
		}
	}

	// Find columns from header
	var dateColumn, payeeColumn, amountColumn int
	dateColumn, payeeColumn, amountColumn = -1, -1, -1
	for fieldIndex, fieldName := range csvRecords[0] {
		fieldName = strings.ToLower(fieldName)
		if strings.Contains(fieldName, "date") {
			dateColumn = fieldIndex
		} else if strings.Contains(fieldName, "description") {
			payeeColumn = fieldIndex
		} else if strings.Contains(fieldName, "payee") {
			payeeColumn = fieldIndex
		} else if strings.Contains(fieldName, "amount") {
			amountColumn = fieldIndex
		} else if strings.Contains(fieldName, "expense") {
			amountColumn = fieldIndex
		}
	}

	if dateColumn < 0 || payeeColumn < 0 || amountColumn < 0 {
		fmt.Println("Unable to find columns required from header field names.")
		return
	}

	expenseAccount := ledger.Account{Name: "unknown:unknown", Balance: new(big.Rat)}
	csvAccount := ledger.Account{Name: matchingAccount, Balance: new(big.Rat)}
	for _, record := range csvRecords[1:] {
		inputPayeeWords := strings.Split(record[payeeColumn], " ")
		csvDate, _ := time.Parse(csvDateFormat, record[dateColumn])
		if !existingTransaction(generalLedger, csvDate, inputPayeeWords[0]) {
			// Classify into expense account
			_, likely, _ := classifier.LogScores(inputPayeeWords)
			if likely >= 0 {
				expenseAccount.Name = string(classifier.Classes[likely])
			}

			// Negate amount if required
			expenseAccount.Balance.SetString(record[amountColumn])
			if negateAmount {
				expenseAccount.Balance.Neg(expenseAccount.Balance)
			}

			// Csv amount is the negative of the expense amount
			csvAccount.Balance.Neg(expenseAccount.Balance)

			// Create valid transaction for print in ledger format
			trans := &ledger.Transaction{Date: csvDate, Payee: record[payeeColumn]}
			trans.AccountChanges = []ledger.Account{csvAccount, expenseAccount}
			PrintTransaction(trans, 80)
		}
	}
}

func existingTransaction(generalLedger []*ledger.Transaction, transDate time.Time, payee string) bool {
	for _, trans := range generalLedger {
		if trans.Date == transDate && strings.HasPrefix(trans.Payee, payee) {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Objdump is a minimal simulation of the GNU objdump tool,
// just enough to support pprof.
//
// Usage:
//	go tool objdump binary start end
//
// Objdump disassembles the binary starting at the start address and
// stopping at the end address. The start and end addresses are program
// counters written in hexadecimal without a leading 0x prefix.
//
// It prints a sequence of stanzas of the form:
//
//	file:line
//	 address: assembly
//	 address: assembly
//	 ...
//
// Each stanza gives the disassembly for a contiguous range of addresses
// all mapped to the same original source file and line number.
//
// The disassembler is missing (golang.org/issue/7452) but will be added
// before the Go 1.3 release.
//
// This tool is intended for use only by pprof; its interface may change or
// it may be deleted entirely in future releases.
package main

import (
	"bufio"
	"debug/elf"
	"debug/gosym"
	"debug/macho"
	"debug/pe"
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"
)

func printUsage(w *os.File) {
	fmt.Fprintf(w, "usage: objdump binary start end\n")
	fmt.Fprintf(w, "disassembles binary from start PC to end PC.\n")
	fmt.Fprintf(w, "start and end are hexadecimal numbers with no 0x prefix.\n")
}

func usage() {
	printUsage(os.Stderr)
	os.Exit(2)
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("objdump: ")

	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 3 {
		usage()
	}

	f, err := os.Open(flag.Arg(0))
	if err != nil {
		log.Fatal(err)
	}

	textStart, textData, symtab, pclntab, err := loadTables(f)
	if err != nil {
		log.Fatalf("reading %s: %v", flag.Arg(0), err)
	}

	pcln := gosym.NewLineTable(pclntab, textStart)
	tab, err := gosym.NewTable(symtab, pcln)
	if err != nil {
		log.Fatalf("reading %s: %v", flag.Arg(0), err)
	}

	start, err := strconv.ParseUint(flag.Arg(1), 0, 64)
	if err != nil {
		log.Fatalf("invalid start PC: %v", err)
	}
	end, err := strconv.ParseUint(flag.Arg(2), 0, 64)
	if err != nil {
		log.Fatalf("invalid end PC: %v", err)
	}

	stdout := bufio.NewWriter(os.Stdout)

	// For now, find spans of same PC/line/fn and
	// emit them as having dummy instructions.
	var (
		spanPC   uint64
		spanFile string
		spanLine int
		spanFn   *gosym.Func
	)

	flush := func(endPC uint64) {
		if spanPC == 0 {
			return
		}
		fmt.Fprintf(stdout, "%s:%d\n", spanFile, spanLine)
		for pc := spanPC; pc < endPC; pc++ {
			// TODO(rsc): Disassemble instructions here.
			if textStart <= pc && pc-textStart < uint64(len(textData)) {
				fmt.Fprintf(stdout, " %x: byte %#x\n", pc, textData[pc-textStart])
			} else {
				fmt.Fprintf(stdout, " %x: ?\n", pc)
			}
		}
		spanPC = 0
	}

	for pc := start; pc < end; pc++ {
		file, line, fn := tab.PCToLine(pc)
		if file != spanFile || line != spanLine || fn != spanFn {
			flush(pc)
			spanPC, spanFile, spanLine, spanFn = pc, file, line, fn
		}
	}
	flush(end)

	stdout.Flush()
}

func loadTables(f *os.File) (textStart uint64, textData, symtab, pclntab []byte, err error) {
	if obj, err := elf.NewFile(f); err == nil {
		if sect := obj.Section(".text"); sect != nil {
			textStart = sect.Addr
			textData, _ = sect.Data()
		}
		if sect := obj.Section(".gosymtab"); sect != nil {
			if symtab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		if sect := obj.Section(".gopclntab"); sect != nil {
			if pclntab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		return textStart, textData, symtab, pclntab, nil
	}

	if obj, err := macho.NewFile(f); err == nil {
		if sect := obj.Section("__text"); sect != nil {
			textStart = sect.Addr
			textData, _ = sect.Data()
		}
		if sect := obj.Section("__gosymtab"); sect != nil {
			if symtab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		if sect := obj.Section("__gopclntab"); sect != nil {
			if pclntab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		return textStart, textData, symtab, pclntab, nil
	}

	if obj, err := pe.NewFile(f); err == nil {
		if sect := obj.Section(".text"); sect != nil {
			textStart = uint64(sect.VirtualAddress)
			textData, _ = sect.Data()
		}
		if sect := obj.Section(".gosymtab"); sect != nil {
			if symtab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		if sect := obj.Section(".gopclntab"); sect != nil {
			if pclntab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		return textStart, textData, symtab, pclntab, nil
	}

	return 0, nil, nil, nil, fmt.Errorf("unrecognized binary format")
}
<commit_msg>cmd/objdump: actually accept hex address without "0x" prefix. Fixes issue 7936.<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Objdump is a minimal simulation of the GNU objdump tool,
// just enough to support pprof.
//
// Usage:
//	go tool objdump binary start end
//
// Objdump disassembles the binary starting at the start address and
// stopping at the end address. The start and end addresses are program
// counters written in hexadecimal with optional leading 0x prefix.
//
// It prints a sequence of stanzas of the form:
//
//	file:line
//	 address: assembly
//	 address: assembly
//	 ...
//
// Each stanza gives the disassembly for a contiguous range of addresses
// all mapped to the same original source file and line number.
//
// The disassembler is missing (golang.org/issue/7452) but will be added
// before the Go 1.3 release.
//
// This tool is intended for use only by pprof; its interface may change or
// it may be deleted entirely in future releases.
package main

import (
	"bufio"
	"debug/elf"
	"debug/gosym"
	"debug/macho"
	"debug/pe"
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
)

func printUsage(w *os.File) {
	fmt.Fprintf(w, "usage: objdump binary start end\n")
	fmt.Fprintf(w, "disassembles binary from start PC to end PC.\n")
	fmt.Fprintf(w, "start and end are hexadecimal numbers with optional leading 0x prefix.\n")
}

func usage() {
	printUsage(os.Stderr)
	os.Exit(2)
}

func main() {
	log.SetFlags(0)
	log.SetPrefix("objdump: ")

	flag.Usage = usage
	flag.Parse()
	if flag.NArg() != 3 {
		usage()
	}

	f, err := os.Open(flag.Arg(0))
	if err != nil {
		log.Fatal(err)
	}

	textStart, textData, symtab, pclntab, err := loadTables(f)
	if err != nil {
		log.Fatalf("reading %s: %v", flag.Arg(0), err)
	}

	pcln := gosym.NewLineTable(pclntab, textStart)
	tab, err := gosym.NewTable(symtab, pcln)
	if err != nil {
		log.Fatalf("reading %s: %v", flag.Arg(0), err)
	}

	start, err := strconv.ParseUint(strings.TrimPrefix(flag.Arg(1), "0x"), 16, 64)
	if err != nil {
		log.Fatalf("invalid start PC: %v", err)
	}
	end, err := strconv.ParseUint(strings.TrimPrefix(flag.Arg(2), "0x"), 16, 64)
	if err != nil {
		log.Fatalf("invalid end PC: %v", err)
	}

	stdout := bufio.NewWriter(os.Stdout)

	// For now, find spans of same PC/line/fn and
	// emit them as having dummy instructions.
	var (
		spanPC   uint64
		spanFile string
		spanLine int
		spanFn   *gosym.Func
	)

	flush := func(endPC uint64) {
		if spanPC == 0 {
			return
		}
		fmt.Fprintf(stdout, "%s:%d\n", spanFile, spanLine)
		for pc := spanPC; pc < endPC; pc++ {
			// TODO(rsc): Disassemble instructions here.
			if textStart <= pc && pc-textStart < uint64(len(textData)) {
				fmt.Fprintf(stdout, " %x: byte %#x\n", pc, textData[pc-textStart])
			} else {
				fmt.Fprintf(stdout, " %x: ?\n", pc)
			}
		}
		spanPC = 0
	}

	for pc := start; pc < end; pc++ {
		file, line, fn := tab.PCToLine(pc)
		if file != spanFile || line != spanLine || fn != spanFn {
			flush(pc)
			spanPC, spanFile, spanLine, spanFn = pc, file, line, fn
		}
	}
	flush(end)

	stdout.Flush()
}

func loadTables(f *os.File) (textStart uint64, textData, symtab, pclntab []byte, err error) {
	if obj, err := elf.NewFile(f); err == nil {
		if sect := obj.Section(".text"); sect != nil {
			textStart = sect.Addr
			textData, _ = sect.Data()
		}
		if sect := obj.Section(".gosymtab"); sect != nil {
			if symtab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		if sect := obj.Section(".gopclntab"); sect != nil {
			if pclntab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		return textStart, textData, symtab, pclntab, nil
	}

	if obj, err := macho.NewFile(f); err == nil {
		if sect := obj.Section("__text"); sect != nil {
			textStart = sect.Addr
			textData, _ = sect.Data()
		}
		if sect := obj.Section("__gosymtab"); sect != nil {
			if symtab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		if sect := obj.Section("__gopclntab"); sect != nil {
			if pclntab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		return textStart, textData, symtab, pclntab, nil
	}

	if obj, err := pe.NewFile(f); err == nil {
		if sect := obj.Section(".text"); sect != nil {
			textStart = uint64(sect.VirtualAddress)
			textData, _ = sect.Data()
		}
		if sect := obj.Section(".gosymtab"); sect != nil {
			if symtab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		if sect := obj.Section(".gopclntab"); sect != nil {
			if pclntab, err = sect.Data(); err != nil {
				return 0, nil, nil, nil, err
			}
		}
		return textStart, textData, symtab, pclntab, nil
	}

	return 0, nil, nil, nil, fmt.Errorf("unrecognized binary format")
}
<|endoftext|>


<commit_before><commit_msg>vsock: skip Dial integration test for CI when applicable<commit_after><|endoftext|>


<commit_before>package main

import (
	"fmt"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildmaster/proto"
	"golang.org/x/net/context"
)

func (s *Server) updateWorld(ctx context.Context, server *pbd.RegistryEntry) ([]string, error) {
	jobs, err := s.getter.getJobs(ctx, server)
	if err != nil {
		return []string{}, err
	}

	slaveMap := []string{}
	for _, job := range jobs {
		slaveMap = append(slaveMap, job.GetJob().GetName())
	}

	return slaveMap, nil
}

func (s *Server) adjustWorld(ctx context.Context) error {
	slaves, err := s.getter.getSlaves()
	if err != nil {
		return err
	}

	var ourSlave *pbd.RegistryEntry
	for _, slave := range slaves.GetServices() {
		if slave.Identifier == s.Registry.Identifier {
			ourSlave = slave
		}
	}
	if ourSlave == nil {
		return fmt.Errorf("Cannot locate local gbs from %v", slaves)
	}

	if len(slaves.GetServices()) == 0 {
		return fmt.Errorf("Unable to locate any slaves")
	}

	jobCount := make(map[string]int)
	ourjobs := make(map[string]bool)
	for _, server := range slaves.GetServices() {
		slaves, err := s.updateWorld(ctx, server)
		if err != nil {
			return err
		}

		for _, j := range slaves {
			jobCount[j]++
			if server.Identifier == s.Registry.Identifier {
				ourjobs[j] = true
			}
		}
	}

	localConfig, err := s.getter.getConfig(ctx, ourSlave)
	if err != nil {
		return err
	}

	for _, intent := range s.config.Nintents {
		if !ourjobs[intent.GetJob().GetName()] {
			allmatch := true
			for _, req := range intent.GetJob().GetRequirements() {
				localmatch := false
				for _, r := range localConfig {
					if r.Category == req.Category && r.Properties == req.Properties {
						localmatch = true
					}
				}

				if !localmatch {
					allmatch = false
				}
			}

			if allmatch {
				err := s.check(ctx, intent, jobCount, ourSlave)
				if err != nil {
					return err
				}
			}
		}
	}

	return nil
}

func (s *Server) check(ctx context.Context, i *pb.NIntent, counts map[string]int, ls *pbd.RegistryEntry) error {
	if i.Redundancy == pb.Redundancy_GLOBAL {
		return s.runJob(ctx, i.GetJob(), ls)
	}

	if i.Redundancy == pb.Redundancy_REDUNDANT {
		if counts[i.GetJob().GetName()] < 3 {
			return s.runJob(ctx, i.GetJob(), ls)
		}
	}

	if counts[i.GetJob().GetName()] < int(i.Count) {
		return s.runJob(ctx, i.GetJob(), ls)
	}

	return nil
}
<commit_msg>Ignore erroring slaves when adjusting world<commit_after>package main

import (
	"fmt"

	pbd "github.com/brotherlogic/discovery/proto"
	pb "github.com/brotherlogic/gobuildmaster/proto"
	"golang.org/x/net/context"
)

func (s *Server) updateWorld(ctx context.Context, server *pbd.RegistryEntry) ([]string, error) {
	jobs, err := s.getter.getJobs(ctx, server)
	if err != nil {
		return []string{}, err
	}

	slaveMap := []string{}
	for _, job := range jobs {
		slaveMap = append(slaveMap, job.GetJob().GetName())
	}

	return slaveMap, nil
}

func (s *Server) adjustWorld(ctx context.Context) error {
	slaves, err := s.getter.getSlaves()
	if err != nil {
		return err
	}

	var ourSlave *pbd.RegistryEntry
	for _, slave := range slaves.GetServices() {
		if slave.Identifier == s.Registry.Identifier {
			ourSlave = slave
		}
	}
	if ourSlave == nil {
		return fmt.Errorf("Cannot locate local gbs from %v", slaves)
	}

	if len(slaves.GetServices()) == 0 {
		return fmt.Errorf("Unable to locate any slaves")
	}

	jobCount := make(map[string]int)
	ourjobs := make(map[string]bool)
	for _, server := range slaves.GetServices() {
		slaves, err := s.updateWorld(ctx, server)
		if err != nil {
			s.Log(fmt.Sprintf("Unable to reach %v -> %v", server, err))
			continue
		}

		for _, j := range slaves {
			jobCount[j]++
			if server.Identifier == s.Registry.Identifier {
				ourjobs[j] = true
			}
		}
	}

	localConfig, err := s.getter.getConfig(ctx, ourSlave)
	if err != nil {
		return err
	}

	for _, intent := range s.config.Nintents {
		if !ourjobs[intent.GetJob().GetName()] {
			allmatch := true
			for _, req := range intent.GetJob().GetRequirements() {
				localmatch := false
				for _, r := range localConfig {
					if r.Category == req.Category && r.Properties == req.Properties {
						localmatch = true
					}
				}

				if !localmatch {
					allmatch = false
				}
			}

			if allmatch {
				err := s.check(ctx, intent, jobCount, ourSlave)
				if err != nil {
					return err
				}
			}
		}
	}

	return nil
}

func (s *Server) check(ctx context.Context, i *pb.NIntent, counts map[string]int, ls *pbd.RegistryEntry) error {
	if i.Redundancy == pb.Redundancy_GLOBAL {
		return s.runJob(ctx, i.GetJob(), ls)
	}

	if i.Redundancy == pb.Redundancy_REDUNDANT {
		if counts[i.GetJob().GetName()] < 3 {
			return s.runJob(ctx, i.GetJob(), ls)
		}
	}

	if counts[i.GetJob().GetName()] < int(i.Count) {
		return s.runJob(ctx, i.GetJob(), ls)
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fusetesting

import (
	"fmt"
	"os"
	"reflect"
	"time"

	"github.com/jacobsa/oglematchers"
)

// Match os.FileInfo values that specify an mtime equal to the given time. On
// platforms where the Sys() method returns a struct containing an mtime, check
// also that it matches.
func MtimeIs(expected time.Time) oglematchers.Matcher {
	return oglematchers.NewMatcher(
		func(c interface{}) error { return mtimeIs(c, expected) },
		fmt.Sprintf("mtime is %v", expected))
}

func mtimeIs(c interface{}, expected time.Time) error {
	fi, ok := c.(os.FileInfo)
	if !ok {
		return fmt.Errorf("which is of type %v", reflect.TypeOf(c))
	}

	// Check ModTime().
	if fi.ModTime() != expected {
		d := fi.ModTime().Sub(expected)
		return fmt.Errorf("which has mtime %v, off by %v", fi.ModTime(), d)
	}

	// Check Sys().
	if sysMtime, ok := extractMtime(fi.Sys()); ok {
		if sysMtime != expected {
			d := sysMtime.Sub(expected)
			return fmt.Errorf("which has Sys() mtime %v, off by %v", sysMtime, d)
		}
	}

	return nil
}

// Extract the mtime from the result of os.FileInfo.Sys(), in a
// platform-specific way. If not supported on this platform, return !ok.
//
// Defined in stat_darwin.go, etc.
func extractMtime(sys interface{}) (mtime time.Time, ok bool)

// Match os.FileInfo values that specify a file birth time equal to the given
// time. On platforms where there is no birth time available, match all
// os.FileInfo values.
func BirthtimeIs(expected time.Time) oglematchers.Matcher
<commit_msg>Fixed a build error.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package fusetesting

import (
	"fmt"
	"os"
	"reflect"
	"time"

	"github.com/jacobsa/oglematchers"
)

// Match os.FileInfo values that specify an mtime equal to the given time. On
// platforms where the Sys() method returns a struct containing an mtime, check
// also that it matches.
func MtimeIs(expected time.Time) oglematchers.Matcher {
	return oglematchers.NewMatcher(
		func(c interface{}) error { return mtimeIs(c, expected) },
		fmt.Sprintf("mtime is %v", expected))
}

func mtimeIs(c interface{}, expected time.Time) error {
	fi, ok := c.(os.FileInfo)
	if !ok {
		return fmt.Errorf("which is of type %v", reflect.TypeOf(c))
	}

	// Check ModTime().
	if fi.ModTime() != expected {
		d := fi.ModTime().Sub(expected)
		return fmt.Errorf("which has mtime %v, off by %v", fi.ModTime(), d)
	}

	// Check Sys().
	if sysMtime, ok := extractMtime(fi.Sys()); ok {
		if sysMtime != expected {
			d := sysMtime.Sub(expected)
			return fmt.Errorf("which has Sys() mtime %v, off by %v", sysMtime, d)
		}
	}

	return nil
}

// Match os.FileInfo values that specify a file birth time equal to the given
// time. On platforms where there is no birth time available, match all
// os.FileInfo values.
func BirthtimeIs(expected time.Time) oglematchers.Matcher
<|endoftext|>


<commit_before>// Copyright 2016 The JoeFriday authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package mem gets and processes /proc/meminfo, returning the data in the
// appropriate format.
package mem

import (
	"bytes"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"time"

	fb "github.com/google/flatbuffers/go"
	joe "github.com/mohae/joefriday"
)

type Info struct {
	Timestamp    int64
	MemTotal     int
	MemFree      int
	MemAvailable int
	Buffers      int
	Cached       int
	SwapCached   int
	Active       int
	Inactive     int
	SwapTotal    int
	SwapFree     int
}

// Serialize serializes the Info using flatbuffers.
func (i *Info) Serialize() []byte {
	bldr := fb.NewBuilder(0)
	DataStart(bldr)
	DataAddTimestamp(bldr, int64(i.Timestamp))
	DataAddMemTotal(bldr, int64(i.MemTotal))
	DataAddMemFree(bldr, int64(i.MemFree))
	DataAddMemAvailable(bldr, int64(i.MemAvailable))
	DataAddBuffers(bldr, int64(i.Buffers))
	DataAddCached(bldr, int64(i.Cached))
	DataAddSwapCached(bldr, int64(i.SwapCached))
	DataAddActive(bldr, int64(i.Active))
	DataAddInactive(bldr, int64(i.Inactive))
	DataAddSwapTotal(bldr, int64(i.SwapTotal))
	DataAddSwapFree(bldr, int64(i.SwapFree))
	bldr.Finish(DataEnd(bldr))
	return bldr.Bytes[bldr.Head():]
}

// Deserialize deserializes bytes representing flatbuffers serialized Data
// into *Info.  If the bytes are not from flatbuffers serialization of
// Data, it is a programmer error and a panic will occur.
func Deserialize(p []byte) *Info {
	data := GetRootAsData(p, 0)
	info := &Info{}
	info.Timestamp = data.Timestamp()
	info.MemTotal = int(data.MemTotal())
	info.MemFree = int(data.MemFree())
	info.MemAvailable = int(data.MemAvailable())
	info.Buffers = int(data.Buffers())
	info.Cached = int(data.Cached())
	info.SwapCached = int(data.SwapCached())
	info.Active = int(data.Active())
	info.Inactive = int(data.Inactive())
	info.SwapTotal = int(data.SwapTotal())
	info.SwapFree = int(data.SwapFree())
	return info
}

// GetInfo returns some of the results of /proc/meminfo.
func GetInfo() (*Info, error) {
	var out bytes.Buffer
	var l, i int
	var name string
	var err error
	var v byte
	t := time.Now().UTC().UnixNano()
	err = meminfo(&out)
	if err != nil {
		return nil, err
	}
	inf := &Info{Timestamp: t}
	var pos int
	line := make([]byte, 0, 50)
	val := make([]byte, 0, 32)
	for {
		if l == 16 {
			break
		}
		line, err = out.ReadBytes(joe.LF)
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, fmt.Errorf("error reading output bytes: %s", err)
		}
		l++
		if l > 8 && l < 15 {
			continue
		}
		// first grab the key name (everything up to the ':')
		for i, v = range line {
			if v == 0x3A {
				pos = i + 1
				break
			}
			val = append(val, v)
		}
		name = string(val[:])
		val = val[:0]
		// skip all spaces
		for i, v = range line[pos:] {
			if v != 0x20 {
				pos += i
				break
			}
		}

		// grab the numbers
		for _, v = range line[pos:] {
			if v == 0x20 || v == joe.LF {
				break
			}
			val = append(val, v)
		}
		// any conversion error results in 0
		i, err = strconv.Atoi(string(val[:]))
		if err != nil {
			return nil, fmt.Errorf("%s: %s", name, err)
		}
		val = val[:0]
		if name == "MemTotal" {
			inf.MemTotal = i
			continue
		}
		if name == "MemFree" {
			inf.MemFree = i
			continue
		}
		if name == "MemAvailable" {
			inf.MemAvailable = i
			continue
		}
		if name == "Buffers" {
			inf.Buffers = i
			continue
		}
		if name == "Cached" {
			inf.MemAvailable = i
			continue
		}
		if name == "SwapCached" {
			inf.SwapCached = i
			continue
		}
		if name == "Active" {
			inf.Active = i
			continue
		}
		if name == "Inactive" {
			inf.Inactive = i
			continue
		}
		if name == "SwapTotal" {
			inf.SwapTotal = i
			continue
		}
		if name == "SwapFree" {
			inf.SwapFree = i
			continue
		}
	}
	return inf, nil
}

// GetData returns the current meminfo as flatbuffer serialized bytes.
// TODO: Benchmark to see if we should just create the flatbuffers w/o
// doing the intermediate step of to the data structure.
func GetData() ([]byte, error) {
	inf, err := GetInfo()
	if err != nil {
		return nil, err
	}
	return inf.Serialize(), nil
}

// DataTicker gathers the meminfo on a ticker, whose interval is defined by
// the received duration, and sends the results to the channel.  The output
// is Flatbuffers serialized Data.  Any error encountered during processing
// is sent to the error channel.  Processing will continue
//
// Either closing the done channel or sending struct{} to the done channel
// will result in function exit.  The out channel is closed on exit.
//
// This pre-allocates the builder and everything other than the []byte that
// gets sent to the out channel to reduce allocations, as this is expected
// to be both a frequent and a long-running process.  Doing so reduces
// byte allocations per tick just ~ 42%.
func DataTicker(interval time.Duration, outCh chan []byte, done chan struct{}, errCh chan error) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	defer close(outCh)
	// predeclare some vars
	var out bytes.Buffer
	var l, i, pos int
	var t int64
	var err error
	var v byte
	var name string
	// premake some temp slices
	line := make([]byte, 0, 50)
	val := make([]byte, 0, 32)
	// just reset the bldr at the end of every ticker
	bldr := fb.NewBuilder(0)

	// ticker
	for {
		select {
		case <-done:
			return
		case <-ticker.C:
			cmd := exec.Command("cat", "/proc/meminfo")
			cmd.Stdout = &out
			// The current timestamp is always in UTC
			t = time.Now().UTC().UnixNano()
			err = cmd.Run()
			if err != nil {
				errCh <- joe.Error{Type: "mem", Op: "cat /proc/meminfo", Err: err}
				continue
			}
			DataStart(bldr)
			DataAddTimestamp(bldr, t)
			for {
				if l == 16 {
					break
				}
				line, err = out.ReadBytes(joe.CR)
				if err != nil {
					if err == io.EOF {
						break
					}
					errCh <- joe.Error{Type: "mem", Op: "read command results", Err: err}
					continue
				}
				l++
				if l > 8 || l < 15 {
					continue
				}
				// first grab the key name (everything up to the ':')
				for i, v = range line {
					if v == 0x3A {
						pos = i + 1
						break
					}
					val = append(val, v)
				}
				name = string(val[:])
				val = val[:0]
				// skip all spaces
				for i, v = range line[pos:] {
					if v != 0x20 {
						pos += i
						break
					}
				}

				// grab the numbers
				for _, v = range line[pos:] {
					if v == 0x20 || v == joe.LF {
						break
					}
					val = append(val, v)
				}
				// any conversion error results in 0
				i, err = strconv.Atoi(string(val[:]))
				if err != nil {
					errCh <- joe.Error{Type: "mem", Op: "convert to int", Err: err}
					continue
				}
				val = val[:0]
				if name == "MemTotal" {
					DataAddMemTotal(bldr, int64(i))
					continue
				}
				if name == "MemFree" {
					DataAddMemFree(bldr, int64(i))
					continue
				}
				if name == "MemAvailable" {
					DataAddMemAvailable(bldr, int64(i))
					continue
				}
				if name == "Buffers" {
					DataAddBuffers(bldr, int64(i))
					continue
				}
				if name == "Cached" {
					DataAddMemAvailable(bldr, int64(i))
					continue
				}
				if name == "SwapCached" {
					DataAddSwapCached(bldr, int64(i))
					continue
				}
				if name == "Active" {
					DataAddActive(bldr, int64(i))
					continue
				}
				if name == "Inactive" {
					DataAddInactive(bldr, int64(i))
					continue
				}
				if name == "SwapTotal" {
					DataAddSwapTotal(bldr, int64(i))
					continue
				}
				if name == "SwapFree" {
					DataAddSwapFree(bldr, int64(i))
					continue
				}
			}
			bldr.Finish(DataEnd(bldr))
			data := bldr.Bytes[bldr.Head():]
			outCh <- data
			bldr.Reset()
			out.Reset()
			l = 0
		}
	}
}

func (d *Data) String() string {
	return fmt.Sprintf("Timestamp: %v\nMemTotal:\t%d\tMemFree:\t%d\tMemAvailable:\t%d\tActive:\t%d\tInactive:\t%d\nCached:\t%d\tBuffers\t:%d\nSwapTotal:\t%d\tSwapCached:\t%d\tSwapFree:\t%d\n", time.Unix(0, d.Timestamp()).UTC(), d.MemTotal(), d.MemFree(), d.MemAvailable(), d.Active(), d.Inactive(), d.Cached(), d.Buffers(), d.SwapTotal(), d.SwapCached(), d.SwapFree())
}

func meminfo(buff *bytes.Buffer) error {
	cmd := exec.Command("cat", "/proc/meminfo")
	cmd.Stdout = buff
	return cmd.Run()
}
<commit_msg>fix bugs in DataTicker parsing logic<commit_after>// Copyright 2016 The JoeFriday authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package mem gets and processes /proc/meminfo, returning the data in the
// appropriate format.
package mem

import (
	"bytes"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"time"

	fb "github.com/google/flatbuffers/go"
	joe "github.com/mohae/joefriday"
)

type Info struct {
	Timestamp    int64
	MemTotal     int
	MemFree      int
	MemAvailable int
	Buffers      int
	Cached       int
	SwapCached   int
	Active       int
	Inactive     int
	SwapTotal    int
	SwapFree     int
}

// Serialize serializes the Info using flatbuffers.
func (i *Info) Serialize() []byte {
	bldr := fb.NewBuilder(0)
	DataStart(bldr)
	DataAddTimestamp(bldr, int64(i.Timestamp))
	DataAddMemTotal(bldr, int64(i.MemTotal))
	DataAddMemFree(bldr, int64(i.MemFree))
	DataAddMemAvailable(bldr, int64(i.MemAvailable))
	DataAddBuffers(bldr, int64(i.Buffers))
	DataAddCached(bldr, int64(i.Cached))
	DataAddSwapCached(bldr, int64(i.SwapCached))
	DataAddActive(bldr, int64(i.Active))
	DataAddInactive(bldr, int64(i.Inactive))
	DataAddSwapTotal(bldr, int64(i.SwapTotal))
	DataAddSwapFree(bldr, int64(i.SwapFree))
	bldr.Finish(DataEnd(bldr))
	return bldr.Bytes[bldr.Head():]
}

// Deserialize deserializes bytes representing flatbuffers serialized Data
// into *Info.  If the bytes are not from flatbuffers serialization of
// Data, it is a programmer error and a panic will occur.
func Deserialize(p []byte) *Info {
	data := GetRootAsData(p, 0)
	info := &Info{}
	info.Timestamp = data.Timestamp()
	info.MemTotal = int(data.MemTotal())
	info.MemFree = int(data.MemFree())
	info.MemAvailable = int(data.MemAvailable())
	info.Buffers = int(data.Buffers())
	info.Cached = int(data.Cached())
	info.SwapCached = int(data.SwapCached())
	info.Active = int(data.Active())
	info.Inactive = int(data.Inactive())
	info.SwapTotal = int(data.SwapTotal())
	info.SwapFree = int(data.SwapFree())
	return info
}

// GetInfo returns some of the results of /proc/meminfo.
func GetInfo() (*Info, error) {
	var out bytes.Buffer
	var l, i int
	var name string
	var err error
	var v byte
	t := time.Now().UTC().UnixNano()
	err = meminfo(&out)
	if err != nil {
		return nil, err
	}
	inf := &Info{Timestamp: t}
	var pos int
	line := make([]byte, 0, 50)
	val := make([]byte, 0, 32)
	for {
		if l == 16 {
			break
		}
		line, err = out.ReadBytes(joe.LF)
		if err != nil {
			if err == io.EOF {
				break
			}
			return nil, fmt.Errorf("error reading output bytes: %s", err)
		}
		l++
		if l > 8 && l < 15 {
			continue
		}
		// first grab the key name (everything up to the ':')
		for i, v = range line {
			if v == 0x3A {
				pos = i + 1
				break
			}
			val = append(val, v)
		}
		name = string(val[:])
		val = val[:0]
		// skip all spaces
		for i, v = range line[pos:] {
			if v != 0x20 {
				pos += i
				break
			}
		}

		// grab the numbers
		for _, v = range line[pos:] {
			if v == 0x20 || v == joe.CR {
				break
			}
			val = append(val, v)
		}
		// any conversion error results in 0
		i, err = strconv.Atoi(string(val[:]))
		if err != nil {
			return nil, fmt.Errorf("%s: %s", name, err)
		}
		val = val[:0]
		if name == "MemTotal" {
			inf.MemTotal = i
			continue
		}
		if name == "MemFree" {
			inf.MemFree = i
			continue
		}
		if name == "MemAvailable" {
			inf.MemAvailable = i
			continue
		}
		if name == "Buffers" {
			inf.Buffers = i
			continue
		}
		if name == "Cached" {
			inf.MemAvailable = i
			continue
		}
		if name == "SwapCached" {
			inf.SwapCached = i
			continue
		}
		if name == "Active" {
			inf.Active = i
			continue
		}
		if name == "Inactive" {
			inf.Inactive = i
			continue
		}
		if name == "SwapTotal" {
			inf.SwapTotal = i
			continue
		}
		if name == "SwapFree" {
			inf.SwapFree = i
			continue
		}
	}
	return inf, nil
}

// GetData returns the current meminfo as flatbuffer serialized bytes.
// TODO: Benchmark to see if we should just create the flatbuffers w/o
// doing the intermediate step of to the data structure.
func GetData() ([]byte, error) {
	inf, err := GetInfo()
	if err != nil {
		return nil, err
	}
	return inf.Serialize(), nil
}

// DataTicker gathers the meminfo on a ticker, whose interval is defined by
// the received duration, and sends the results to the channel.  The output
// is Flatbuffers serialized Data.  Any error encountered during processing
// is sent to the error channel.  Processing will continue
//
// Either closing the done channel or sending struct{} to the done channel
// will result in function exit.  The out channel is closed on exit.
//
// This pre-allocates the builder and everything other than the []byte that
// gets sent to the out channel to reduce allocations, as this is expected
// to be both a frequent and a long-running process.  Doing so reduces
// byte allocations per tick just ~ 42%.
func DataTicker(interval time.Duration, outCh chan []byte, done chan struct{}, errCh chan error) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	defer close(outCh)
	// predeclare some vars
	var out bytes.Buffer
	var l, i, pos int
	var t int64
	var err error
	var v byte
	var name string
	// premake some temp slices
	line := make([]byte, 0, 50)
	val := make([]byte, 0, 32)
	// just reset the bldr at the end of every ticker
	bldr := fb.NewBuilder(0)

	// ticker
	for {
		select {
		case <-done:
			return
		case <-ticker.C:
			cmd := exec.Command("cat", "/proc/meminfo")
			cmd.Stdout = &out
			// The current timestamp is always in UTC
			t = time.Now().UTC().UnixNano()
			err = cmd.Run()
			if err != nil {
				errCh <- joe.Error{Type: "mem", Op: "cat /proc/meminfo", Err: err}
				continue
			}
			DataStart(bldr)
			DataAddTimestamp(bldr, t)
			for {
				if l == 16 {
					break
				}
				line, err = out.ReadBytes(joe.LF)
				if err != nil {
					if err == io.EOF {
						break
					}
					errCh <- joe.Error{Type: "mem", Op: "read command results", Err: err}
					continue
				}
				l++
				if l > 8 && l < 15 {
					continue
				}
				// first grab the key name (everything up to the ':')
				for i, v = range line {
					if v == 0x3A {
						pos = i + 1
						break
					}
					val = append(val, v)
				}
				name = string(val[:])
				val = val[:0]
				// skip all spaces
				for i, v = range line[pos:] {
					if v != 0x20 {
						pos += i
						break
					}
				}

				// grab the numbers
				for _, v = range line[pos:] {
					if v == 0x20 || v == joe.CR {
						break
					}
					val = append(val, v)
				}
				// any conversion error results in 0
				i, err = strconv.Atoi(string(val[:]))
				if err != nil {
					errCh <- joe.Error{Type: "mem", Op: "convert to int", Err: err}
					continue
				}
				val = val[:0]
				if name == "MemTotal" {
					DataAddMemTotal(bldr, int64(i))
					continue
				}
				if name == "MemFree" {
					DataAddMemFree(bldr, int64(i))
					continue
				}
				if name == "MemAvailable" {
					DataAddMemAvailable(bldr, int64(i))
					continue
				}
				if name == "Buffers" {
					DataAddBuffers(bldr, int64(i))
					continue
				}
				if name == "Cached" {
					DataAddMemAvailable(bldr, int64(i))
					continue
				}
				if name == "SwapCached" {
					DataAddSwapCached(bldr, int64(i))
					continue
				}
				if name == "Active" {
					DataAddActive(bldr, int64(i))
					continue
				}
				if name == "Inactive" {
					DataAddInactive(bldr, int64(i))
					continue
				}
				if name == "SwapTotal" {
					DataAddSwapTotal(bldr, int64(i))
					continue
				}
				if name == "SwapFree" {
					DataAddSwapFree(bldr, int64(i))
					continue
				}
			}
			bldr.Finish(DataEnd(bldr))
			data := bldr.Bytes[bldr.Head():]
			outCh <- data
			bldr.Reset()
			out.Reset()
			l = 0
		}
	}
}

func (d *Data) String() string {
	return fmt.Sprintf("Timestamp: %v\nMemTotal:\t%d\tMemFree:\t%d\tMemAvailable:\t%d\tActive:\t%d\tInactive:\t%d\nCached:\t\t%d\tBuffers\t:%d\nSwapTotal:\t%d\tSwapCached:\t%d\tSwapFree:\t%d\n", time.Unix(0, d.Timestamp()).UTC(), d.MemTotal(), d.MemFree(), d.MemAvailable(), d.Active(), d.Inactive(), d.Cached(), d.Buffers(), d.SwapTotal(), d.SwapCached(), d.SwapFree())
}

func meminfo(buff *bytes.Buffer) error {
	cmd := exec.Command("cat", "/proc/meminfo")
	cmd.Stdout = buff
	return cmd.Run()
}
<|endoftext|>


<commit_before>// Copyright 2017 gf Author(https://gitee.com/johng/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/johng/gf.

// 文件管理.
package gfile

import (
    "bytes"
    "errors"
    "fmt"
    "gitee.com/johng/gf/g/container/gtype"
    "gitee.com/johng/gf/g/util/gconv"
    "gitee.com/johng/gf/g/util/gregex"
    "io"
    "os"
    "os/exec"
    "os/user"
    "path/filepath"
    "runtime"
    "sort"
    "strings"
    "time"
)

// 文件分隔符
const (
    Separator     = string(filepath.Separator)
    // 默认的文件打开权限
    gDEFAULT_PERM = 0666
)

var (
    // 源码的main包所在目录，仅仅会设置一次
    mainPkgPath   = gtype.NewString()

    // 编译时的 GOROOT 数值
    goRootOfBuild = gtype.NewString()
)

// 给定文件的绝对路径创建文件
func Mkdir(path string) error {
    err  := os.MkdirAll(path, os.ModePerm)
    if err != nil {
        return err
    }
    return nil
}

// 给定文件的绝对路径创建文件
func Create(path string) error {
    dir := Dir(path)
    if !Exists(dir) {
        Mkdir(dir)
    }
    f, err  := os.Create(path)
    if err != nil {
        return err
    }
    f.Close()
    return nil
}

// 打开文件(os.O_RDWR|os.O_CREATE, 0666)
func Open(path string) (*os.File, error) {
    f, err  := os.OpenFile(path, os.O_RDWR|os.O_CREATE, gDEFAULT_PERM)
    if err != nil {
        return nil, err
    }
    return f, nil
}

// 打开文件(带flag)
func OpenWithFlag(path string, flag int) (*os.File, error) {
    f, err  := os.OpenFile(path, flag, gDEFAULT_PERM)
    if err != nil {
        return nil, err
    }
    return f, nil
}

// 打开文件(带flag&perm)
func OpenWithFlagPerm(path string, flag int, perm int) (*os.File, error) {
    f, err  := os.OpenFile(path, flag, os.FileMode(perm))
    if err != nil {
        return nil, err
    }
    return f, nil
}

// 判断所给路径文件/文件夹是否存在
func Exists(path string) bool {
    if _, err := os.Stat(path); !os.IsNotExist(err) {
        return true
    }
    return false
}

// 判断所给路径是否为文件夹
func IsDir(path string) bool {
    s, err := os.Stat(path)
    if err != nil {
        return false
    }
    return s.IsDir()
}

// 获取当前工作目录
func Pwd() string {
    pwd, _ := filepath.Abs(filepath.Dir(os.Args[0]))
    return pwd
}

// 判断所给路径是否为文件
func IsFile(path string) bool {
    s, err := os.Stat(path)
    if err != nil {
        return false
    }
    return !s.IsDir()
}

// 获取文件或目录信息
func Info(path string) *os.FileInfo {
    info, err := os.Stat(path)
    if err != nil {
        return nil
    }
    return &info
}

// 文件移动/重命名
func Move(src string, dst string) error {
    return os.Rename(src, dst)
}

// 文件移动/重命名
func Rename(src string, dst string) error {
    return Move(src, dst)
}

// 文件复制
func Copy(src string, dst string) error {
    srcFile, err := os.Open(src)
    if err != nil {
        return err
    }
    dstFile, err := os.Create(dst)
    if err != nil {
        return err
    }
    _, err = io.Copy(dstFile, srcFile)
    if err != nil {
        return err
    }
    err = dstFile.Sync()
    if err != nil {
        return err
    }
    srcFile.Close()
    dstFile.Close()
    return nil
}

// 返回目录下的文件名列表
func DirNames(path string) ([]string, error) {
    f, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    list, err := f.Readdirnames(-1)
    f.Close()
    if err != nil {
        return nil, err
    }
    return list, nil
}

// 文件名正则匹配查找，第二个可选参数指定返回的列表是否仅为文件名(非绝对路径)，默认返回绝对路径
func Glob(pattern string, onlyNames...bool) ([]string, error) {
    if list, err := filepath.Glob(pattern); err == nil {
        if len(onlyNames) > 0 && onlyNames[0] && len(list) > 0 {
            array := make([]string, len(list))
            for k, v := range list {
                array[k] = Basename(v)
            }
            return array, nil
        }
        return list, nil
    } else {
        return nil, err
    }
}

// 文件/目录删除
func Remove(path string) error {
    return os.RemoveAll(path)
}

// 文件是否可读
func IsReadable(path string) bool {
    result    := true
    file, err := os.OpenFile(path, os.O_RDONLY, gDEFAULT_PERM)
    if err != nil {
        result = false
    }
    file.Close()
    return result
}

// 文件是否可写
func IsWritable(path string) bool {
    result := true
    if IsDir(path) {
        // 如果是目录，那么创建一个临时文件进行写入测试
        tfile := strings.TrimRight(path, Separator) + Separator + gconv.String(time.Now().UnixNano())
        err   := Create(tfile)
        if err != nil || !Exists(tfile){
            result = false
        } else {
            Remove(tfile)
        }
    } else {
        // 如果是文件，那么判断文件是否可打开
        file, err := os.OpenFile(path, os.O_WRONLY, gDEFAULT_PERM)
        if err != nil {
            result = false
        }
        file.Close()
    }
    return result
}

// 修改文件/目录权限
func Chmod(path string, mode os.FileMode) error {
    return os.Chmod(path, mode)
}

// 打开目录，并返回其下一级文件列表(绝对路径)，按照文件名称大小写进行排序，支持目录递归遍历。
func ScanDir(path string, pattern string, recursive ... bool) ([]string, error) {
    list, err := doScanDir(path, pattern, recursive...)
    if err != nil {
        return nil, err
    }
    if len(list) > 0 {
        sort.Strings(list)
    }
    return list, nil
}

// 内部检索目录方法，支持递归，返回没有排序的文件绝对路径列表结果
func doScanDir(path string, pattern string, recursive ... bool) ([]string, error) {
    var list []string
    // 打开目录
    dfile, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    defer dfile.Close()
    // 读取目录下的文件列表
    names, err := dfile.Readdirnames(-1)
    if err != nil {
        return nil, err
    }
    // 是否递归遍历
    for _, name := range names {
        path := fmt.Sprintf("%s%s%s", path, Separator, name)
        if IsDir(path) && len(recursive) > 0 && recursive[0] {
            array, _ := doScanDir(path, pattern, true)
            if len(array) > 0 {
                list = append(list, array...)
            }
        }
        // 满足pattern才加入结果列表
        if match, err := filepath.Match(pattern, name); err == nil && match {
            list = append(list, path)
        }
    }
    return list, nil
}

// 将所给定的路径转换为绝对路径
// 并判断文件路径是否存在，如果文件不存在，那么返回空字符串
func RealPath(path string) string {
    p, err := filepath.Abs(path)
    if err != nil {
        return ""
    }
    if !Exists(p) {
        return ""
    }
    return p
}

// 获取当前执行文件的绝对路径
func SelfPath() string {
    p, _ := filepath.Abs(os.Args[0])
    return p
}

// 获取当前执行文件的目录绝对路径
func SelfDir() string {
    return filepath.Dir(SelfPath())
}

// 获取指定文件路径的文件名称
func Basename(path string) string {
    return filepath.Base(path)
}

// 获取指定文件路径的目录地址绝对路径
func Dir(path string) string {
    return filepath.Dir(path)
}

// 获取指定文件路径的文件扩展名
func Ext(path string) string {
    return filepath.Ext(path)
}

// 获取用户主目录
func Home() (string, error) {
    u, err := user.Current()
    if nil == err {
        return u.HomeDir, nil
    }
    if "windows" == runtime.GOOS {
        return homeWindows()
    }
    return homeUnix()
}

func homeUnix() (string, error) {
    if home := os.Getenv("HOME"); home != "" {
        return home, nil
    }
    var stdout bytes.Buffer
    cmd := exec.Command("sh", "-c", "eval echo ~$USER")
    cmd.Stdout = &stdout
    if err := cmd.Run(); err != nil {
        return "", err
    }

    result := strings.TrimSpace(stdout.String())
    if result == "" {
        return "", errors.New("blank output when reading home directory")
    }

    return result, nil
}

func homeWindows() (string, error) {
    drive := os.Getenv("HOMEDRIVE")
    path  := os.Getenv("HOMEPATH")
    home  := drive + path
    if drive == "" || path == "" {
        home = os.Getenv("USERPROFILE")
    }
    if home == "" {
        return "", errors.New("HOMEDRIVE, HOMEPATH, and USERPROFILE are blank")
    }

    return home, nil
}

// 获取入口函数文件所在目录(main包文件目录)，
// **仅对源码开发环境有效(即仅对生成该可执行文件的系统下有效)**
func MainPkgPath() string {
    path := mainPkgPath.Val()
    if path != "" {
        return path
    }
    f := ""
    for i := 1; i < 10000; i++ {
        if _, file, _, ok := runtime.Caller(i); ok {
            if strings.EqualFold("<autogenerated>", file) {
                // 如果是通过init包方法进入，那么无法得到准确的文件路径
                f = ""
            } else {
                goroot := GoRootOfBuild()
                if goroot != "" && !gregex.IsMatchString("^" + GoRootOfBuild(), file) {
                    // 不包含go源码路径
                    f = file
                }
            }
        } else {
            break
        }
    }
    if f != "" {
       p := Dir(f)
       mainPkgPath.Set(p)
       return p
    }
    return ""
}

// 编译时环境的GOROOT数值(对init初始化方法调用时无效，获取不了ROOT值)
// 注意：可能返回空
func GoRootOfBuild() string {
    if v := goRootOfBuild.Val(); v != "" {
        return v
    }
    firstEntry := ""
    for i := 0; i < 10000; i++ {
        if _, file, _, ok := runtime.Caller(i); ok {
            firstEntry = file
        } else {
            break
        }
    }
    if len(firstEntry) > 0 {
        sep   := "/"
        array := strings.Split(firstEntry, sep)
        if len(array) == 1 {
            sep   = "\\"
            array = strings.Split(firstEntry, sep)
        }
        root := strings.Join(array[0 : len(array) - 3], sep)
        goRootOfBuild.Set(root)
        return root
    }
    return ""
}

// 系统临时目录
func TempDir() string {
    return os.TempDir()
}<commit_msg>完善gfile注释<commit_after>// Copyright 2017 gf Author(https://gitee.com/johng/gf). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://gitee.com/johng/gf.

// 文件管理.
package gfile

import (
    "bytes"
    "errors"
    "fmt"
    "gitee.com/johng/gf/g/container/gtype"
    "gitee.com/johng/gf/g/util/gconv"
    "gitee.com/johng/gf/g/util/gregex"
    "io"
    "os"
    "os/exec"
    "os/user"
    "path/filepath"
    "runtime"
    "sort"
    "strings"
    "time"
)

// 文件分隔符
const (
    Separator     = string(filepath.Separator)
    // 默认的文件打开权限
    gDEFAULT_PERM = 0666
)

var (
    // 源码的main包所在目录，仅仅会设置一次
    mainPkgPath   = gtype.NewString()

    // 编译时的 GOROOT 数值
    goRootOfBuild = gtype.NewString()
)

// 给定文件的绝对路径创建文件
func Mkdir(path string) error {
    err  := os.MkdirAll(path, os.ModePerm)
    if err != nil {
        return err
    }
    return nil
}

// 给定文件的绝对路径创建文件
func Create(path string) error {
    dir := Dir(path)
    if !Exists(dir) {
        Mkdir(dir)
    }
    f, err  := os.Create(path)
    if err != nil {
        return err
    }
    f.Close()
    return nil
}

// 打开文件(os.O_RDWR|os.O_CREATE, 0666)
func Open(path string) (*os.File, error) {
    f, err  := os.OpenFile(path, os.O_RDWR|os.O_CREATE, gDEFAULT_PERM)
    if err != nil {
        return nil, err
    }
    return f, nil
}

// 打开文件(带flag)
func OpenWithFlag(path string, flag int) (*os.File, error) {
    f, err  := os.OpenFile(path, flag, gDEFAULT_PERM)
    if err != nil {
        return nil, err
    }
    return f, nil
}

// 打开文件(带flag&perm)
func OpenWithFlagPerm(path string, flag int, perm int) (*os.File, error) {
    f, err  := os.OpenFile(path, flag, os.FileMode(perm))
    if err != nil {
        return nil, err
    }
    return f, nil
}

// 判断所给路径文件/文件夹是否存在
func Exists(path string) bool {
    if _, err := os.Stat(path); !os.IsNotExist(err) {
        return true
    }
    return false
}

// 判断所给路径是否为文件夹
func IsDir(path string) bool {
    s, err := os.Stat(path)
    if err != nil {
        return false
    }
    return s.IsDir()
}

// 获取当前工作目录
func Pwd() string {
    pwd, _ := filepath.Abs(filepath.Dir(os.Args[0]))
    return pwd
}

// 判断所给路径是否为文件
func IsFile(path string) bool {
    s, err := os.Stat(path)
    if err != nil {
        return false
    }
    return !s.IsDir()
}

// 获取文件或目录信息
func Info(path string) *os.FileInfo {
    info, err := os.Stat(path)
    if err != nil {
        return nil
    }
    return &info
}

// 文件移动/重命名
func Move(src string, dst string) error {
    return os.Rename(src, dst)
}

// 文件移动/重命名
func Rename(src string, dst string) error {
    return Move(src, dst)
}

// 文件复制
func Copy(src string, dst string) error {
    srcFile, err := os.Open(src)
    if err != nil {
        return err
    }
    dstFile, err := os.Create(dst)
    if err != nil {
        return err
    }
    _, err = io.Copy(dstFile, srcFile)
    if err != nil {
        return err
    }
    err = dstFile.Sync()
    if err != nil {
        return err
    }
    srcFile.Close()
    dstFile.Close()
    return nil
}

// 返回目录下的文件名列表
func DirNames(path string) ([]string, error) {
    f, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    list, err := f.Readdirnames(-1)
    f.Close()
    if err != nil {
        return nil, err
    }
    return list, nil
}

// 文件名正则匹配查找，第二个可选参数指定返回的列表是否仅为文件名(非绝对路径)，默认返回绝对路径
func Glob(pattern string, onlyNames...bool) ([]string, error) {
    if list, err := filepath.Glob(pattern); err == nil {
        if len(onlyNames) > 0 && onlyNames[0] && len(list) > 0 {
            array := make([]string, len(list))
            for k, v := range list {
                array[k] = Basename(v)
            }
            return array, nil
        }
        return list, nil
    } else {
        return nil, err
    }
}

// 文件/目录删除
func Remove(path string) error {
    return os.RemoveAll(path)
}

// 文件是否可读
func IsReadable(path string) bool {
    result    := true
    file, err := os.OpenFile(path, os.O_RDONLY, gDEFAULT_PERM)
    if err != nil {
        result = false
    }
    file.Close()
    return result
}

// 文件是否可写
func IsWritable(path string) bool {
    result := true
    if IsDir(path) {
        // 如果是目录，那么创建一个临时文件进行写入测试
        tfile := strings.TrimRight(path, Separator) + Separator + gconv.String(time.Now().UnixNano())
        err   := Create(tfile)
        if err != nil || !Exists(tfile){
            result = false
        } else {
            Remove(tfile)
        }
    } else {
        // 如果是文件，那么判断文件是否可打开
        file, err := os.OpenFile(path, os.O_WRONLY, gDEFAULT_PERM)
        if err != nil {
            result = false
        }
        file.Close()
    }
    return result
}

// 修改文件/目录权限
func Chmod(path string, mode os.FileMode) error {
    return os.Chmod(path, mode)
}

// 打开目录，并返回其下一级文件列表(绝对路径)，按照文件名称大小写进行排序，支持目录递归遍历。
func ScanDir(path string, pattern string, recursive ... bool) ([]string, error) {
    list, err := doScanDir(path, pattern, recursive...)
    if err != nil {
        return nil, err
    }
    if len(list) > 0 {
        sort.Strings(list)
    }
    return list, nil
}

// 内部检索目录方法，支持递归，返回没有排序的文件绝对路径列表结果
func doScanDir(path string, pattern string, recursive ... bool) ([]string, error) {
    var list []string
    // 打开目录
    dfile, err := os.Open(path)
    if err != nil {
        return nil, err
    }
    defer dfile.Close()
    // 读取目录下的文件列表
    names, err := dfile.Readdirnames(-1)
    if err != nil {
        return nil, err
    }
    // 是否递归遍历
    for _, name := range names {
        path := fmt.Sprintf("%s%s%s", path, Separator, name)
        if IsDir(path) && len(recursive) > 0 && recursive[0] {
            array, _ := doScanDir(path, pattern, true)
            if len(array) > 0 {
                list = append(list, array...)
            }
        }
        // 满足pattern才加入结果列表
        if match, err := filepath.Match(pattern, name); err == nil && match {
            list = append(list, path)
        }
    }
    return list, nil
}

// 将所给定的路径转换为绝对路径
// 并判断文件路径是否存在，如果文件不存在，那么返回空字符串
func RealPath(path string) string {
    p, err := filepath.Abs(path)
    if err != nil {
        return ""
    }
    if !Exists(p) {
        return ""
    }
    return p
}

// 获取当前执行文件的绝对路径
func SelfPath() string {
    p, _ := filepath.Abs(os.Args[0])
    return p
}

// 获取当前执行文件的目录绝对路径
func SelfDir() string {
    return filepath.Dir(SelfPath())
}

// 获取指定文件路径的文件名称
func Basename(path string) string {
    return filepath.Base(path)
}

// 获取指定文件路径的目录地址绝对路径
func Dir(path string) string {
    return filepath.Dir(path)
}

// 获取指定文件路径的文件扩展名(包含"."号)
func Ext(path string) string {
    return filepath.Ext(path)
}

// 获取用户主目录
func Home() (string, error) {
    u, err := user.Current()
    if nil == err {
        return u.HomeDir, nil
    }
    if "windows" == runtime.GOOS {
        return homeWindows()
    }
    return homeUnix()
}

func homeUnix() (string, error) {
    if home := os.Getenv("HOME"); home != "" {
        return home, nil
    }
    var stdout bytes.Buffer
    cmd := exec.Command("sh", "-c", "eval echo ~$USER")
    cmd.Stdout = &stdout
    if err := cmd.Run(); err != nil {
        return "", err
    }

    result := strings.TrimSpace(stdout.String())
    if result == "" {
        return "", errors.New("blank output when reading home directory")
    }

    return result, nil
}

func homeWindows() (string, error) {
    drive := os.Getenv("HOMEDRIVE")
    path  := os.Getenv("HOMEPATH")
    home  := drive + path
    if drive == "" || path == "" {
        home = os.Getenv("USERPROFILE")
    }
    if home == "" {
        return "", errors.New("HOMEDRIVE, HOMEPATH, and USERPROFILE are blank")
    }

    return home, nil
}

// 获取入口函数文件所在目录(main包文件目录)，
// **仅对源码开发环境有效(即仅对生成该可执行文件的系统下有效)**
func MainPkgPath() string {
    path := mainPkgPath.Val()
    if path != "" {
        return path
    }
    f := ""
    for i := 1; i < 10000; i++ {
        if _, file, _, ok := runtime.Caller(i); ok {
            if strings.EqualFold("<autogenerated>", file) {
                // 如果是通过init包方法进入，那么无法得到准确的文件路径
                f = ""
            } else {
                goroot := GoRootOfBuild()
                if goroot != "" && !gregex.IsMatchString("^" + GoRootOfBuild(), file) {
                    // 不包含go源码路径
                    f = file
                }
            }
        } else {
            break
        }
    }
    if f != "" {
       p := Dir(f)
       mainPkgPath.Set(p)
       return p
    }
    return ""
}

// 编译时环境的GOROOT数值(对init初始化方法调用时无效，获取不了ROOT值)
// 注意：可能返回空
func GoRootOfBuild() string {
    if v := goRootOfBuild.Val(); v != "" {
        return v
    }
    firstEntry := ""
    for i := 0; i < 10000; i++ {
        if _, file, _, ok := runtime.Caller(i); ok {
            firstEntry = file
        } else {
            break
        }
    }
    if len(firstEntry) > 0 {
        sep   := "/"
        array := strings.Split(firstEntry, sep)
        if len(array) == 1 {
            sep   = "\\"
            array = strings.Split(firstEntry, sep)
        }
        root := strings.Join(array[0 : len(array) - 3], sep)
        goRootOfBuild.Set(root)
        return root
    }
    return ""
}

// 系统临时目录
func TempDir() string {
    return os.TempDir()
}<|endoftext|>


<commit_before>// Package metric provides functions for measuring the difference between the
// values predicted by a model and the values actually observed.
package metric

import (
	"math"
)

// MSE computes the mean-square error.
//
// https://en.wikipedia.org/wiki/Mean_squared_error
func MSE(observations, predictions []float64) float64 {
	var sum, Δ float64

	for i := range observations {
		Δ = predictions[i] - observations[i]
		sum += Δ * Δ
	}

	return sum / float64(len(observations))
}

// RMSE computes the root-mean-square error.
//
// https://en.wikipedia.org/wiki/Root-mean-square_deviation
func RMSE(observations, predictions []float64) float64 {
	return math.Sqrt(MSE(observations, predictions))
}

// NRMSE computes the normalized root-mean-square error.
//
// https://en.wikipedia.org/wiki/Root-mean-square_deviation#Normalized_root-mean-square_deviation
func NRMSE(observations, predictions []float64) float64 {
	count := len(observations)
	if count == 0 {
		return 0
	}

	min, max := observations[0], observations[0]
	for i := 1; i < count; i++ {
		if observations[i] < min {
			min = observations[i]
		}
		if observations[i] > max {
			max = observations[i]
		}
	}

	return RMSE(observations, predictions) / (max - min)
}
<commit_msg>metric: remove an excessive check<commit_after>// Package metric provides functions for measuring the difference between the
// values predicted by a model and the values actually observed.
package metric

import (
	"math"
)

// MSE computes the mean-square error.
//
// https://en.wikipedia.org/wiki/Mean_squared_error
func MSE(observations, predictions []float64) float64 {
	var sum, Δ float64

	for i := range observations {
		Δ = predictions[i] - observations[i]
		sum += Δ * Δ
	}

	return sum / float64(len(observations))
}

// RMSE computes the root-mean-square error.
//
// https://en.wikipedia.org/wiki/Root-mean-square_deviation
func RMSE(observations, predictions []float64) float64 {
	return math.Sqrt(MSE(observations, predictions))
}

// NRMSE computes the normalized root-mean-square error.
//
// https://en.wikipedia.org/wiki/Root-mean-square_deviation#Normalized_root-mean-square_deviation
func NRMSE(observations, predictions []float64) float64 {
	count := len(observations)

	min, max := observations[0], observations[0]
	for i := 1; i < count; i++ {
		if observations[i] < min {
			min = observations[i]
		}
		if observations[i] > max {
			max = observations[i]
		}
	}

	return RMSE(observations, predictions) / (max - min)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

// Package crier reports finished prowjob status to git providers.
package crier

import (
	"context"
	"fmt"
	"time"

	"github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/util/retry"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	ctrlruntimeclient "sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"

	prowv1 "k8s.io/test-infra/prow/apis/prowjobs/v1"
)

type ReportClient interface {
	// Report reports a Prowjob. The provided logger is already populated with the
	// prowjob name and the reporter name.
	// If a reporter wants to defer reporting, it can return a reconcile.Result with a RequeueAfter
	Report(log *logrus.Entry, pj *prowv1.ProwJob) ([]*prowv1.ProwJob, *reconcile.Result, error)
	GetName() string
	// ShouldReport determines if a ProwJob should be reported. The provided logger
	// is already populated with the prowjob name and the reporter name.
	ShouldReport(log *logrus.Entry, pj *prowv1.ProwJob) bool
}

// reconciler struct defines how a controller should encapsulate
// logging, client connectivity, informing (list and watching)
// queueing, and handling of resource changes
type reconciler struct {
	pjclientset       ctrlruntimeclient.Client
	reporter          ReportClient
	enablementChecker func(org, repo string) bool
}

// New constructs a new instance of the crier reconciler.
func New(
	mgr manager.Manager,
	reporter ReportClient,
	numWorkers int,
	enablementChecker func(org, repo string) bool,
) error {
	if err := builder.
		ControllerManagedBy(mgr).
		// Is used for metrics, hence must be unique per controller instance
		Named(fmt.Sprintf("crier_%s", reporter.GetName())).
		For(&prowv1.ProwJob{}).
		WithOptions(controller.Options{MaxConcurrentReconciles: numWorkers}).
		Complete(&reconciler{
			pjclientset:       mgr.GetClient(),
			reporter:          reporter,
			enablementChecker: enablementChecker,
		}); err != nil {
		return fmt.Errorf("failed to construct controller: %w", err)
	}

	return nil
}

func (r *reconciler) updateReportState(ctx context.Context, pj *prowv1.ProwJob, log *logrus.Entry, reportedState prowv1.ProwJobState) error {
	// update pj report status
	newpj := pj.DeepCopy()
	// we set omitempty on PrevReportStates, so here we need to init it if is nil
	if newpj.Status.PrevReportStates == nil {
		newpj.Status.PrevReportStates = map[string]prowv1.ProwJobState{}
	}
	newpj.Status.PrevReportStates[r.reporter.GetName()] = reportedState

	if err := r.pjclientset.Patch(ctx, newpj, ctrlruntimeclient.MergeFrom(pj)); err != nil {
		return fmt.Errorf("failed to patch: %w", err)
	}

	// Block until the update is in the lister to make sure that events from another controller
	// that also does reporting dont trigger another report because our lister doesn't yet contain
	// the updated Status
	name := types.NamespacedName{Namespace: pj.Namespace, Name: pj.Name}
	if err := wait.Poll(100*time.Millisecond, 3*time.Second, func() (bool, error) {
		if err := r.pjclientset.Get(ctx, name, pj); err != nil {
			return false, err
		}
		if pj.Status.PrevReportStates != nil &&
			pj.Status.PrevReportStates[r.reporter.GetName()] == reportedState {
			return true, nil
		}
		return false, nil
	}); err != nil {
		return fmt.Errorf("failed to wait for updated report status to be in lister: %w", err)
	}
	return nil
}

func (r *reconciler) updateReportStateWithRetries(ctx context.Context, pj *prowv1.ProwJob, log *logrus.Entry) error {
	reportState := pj.Status.State
	log = log.WithFields(logrus.Fields{
		"prowjob":   pj.Name,
		"jobName":   pj.Spec.Job,
		"jobStatus": reportState,
	})
	// We have to retry here, if we return we lose the information that we already reported this job.
	if err := retry.RetryOnConflict(retry.DefaultBackoff, func() error {
		// Get it first, this is very cheap
		name := types.NamespacedName{Namespace: pj.Namespace, Name: pj.Name}
		if err := r.pjclientset.Get(ctx, name, pj); err != nil {
			return err
		}
		// Must not wrap until we have kube 1.19, otherwise the RetryOnConflict won't recognize conflicts
		// correctly
		return r.updateReportState(ctx, pj, log, reportState)
	}); err != nil {
		// Very subpar, we will report again. But even if we didn't do that now, we would do so
		// latest when crier gets restarted. In an ideal world, all reporters are idempotent and
		// reporting has no cost.
		return fmt.Errorf("failed to update report state on prowjob: %w", err)
	}

	log.Info("Successfully updated report state on prowjob")
	return nil
}

// Reconcile retrieves each queued item and takes the necessary handler action based off of if
// the item was created or deleted.
func (r *reconciler) Reconcile(ctx context.Context, req reconcile.Request) (reconcile.Result, error) {
	log := logrus.WithField("reporter", r.reporter.GetName()).WithField("key", req.String()).WithField("prowjob", req.Name)
	log.Debug("processing next key")
	result, err := r.reconcile(ctx, log, req)
	if err != nil {
		log.WithError(err).Error("Reconciliation failed")
	}
	if result == nil {
		result = &reconcile.Result{}
	}
	return *result, err
}

func (r *reconciler) reconcile(ctx context.Context, log *logrus.Entry, req reconcile.Request) (*reconcile.Result, error) {

	var pj prowv1.ProwJob
	if err := r.pjclientset.Get(ctx, req.NamespacedName, &pj); err != nil {
		if errors.IsNotFound(err) {
			log.Debug("object no longer exist")
			return nil, nil
		}

		return nil, fmt.Errorf("failed to get prowjob %s: %w", req.String(), err)
	}

	if !r.shouldHandle(&pj) {
		return nil, nil
	}

	log = log.WithField("jobName", pj.Spec.Job)

	if !pj.Spec.Report || !r.reporter.ShouldReport(log, &pj) {
		return nil, nil
	}

	// we set omitempty on PrevReportStates, so here we need to init it if is nil
	if pj.Status.PrevReportStates == nil {
		pj.Status.PrevReportStates = map[string]prowv1.ProwJobState{}
	}

	// already reported current state
	if pj.Status.PrevReportStates[r.reporter.GetName()] == pj.Status.State {
		log.Trace("Already reported")
		return nil, nil
	}

	log = log.WithField("jobStatus", pj.Status.State)
	log.Info("Will report state")
	pjs, requeue, err := r.reporter.Report(log, &pj)
	if err != nil {
		log.WithError(err).Error("failed to report job")
		return nil, fmt.Errorf("failed to report job: %w", err)
	}
	if requeue != nil {
		return requeue, nil
	}

	log.Info("Reported job(s), now will update pj(s)")
	for _, pjob := range pjs {
		if err := r.updateReportStateWithRetries(ctx, pjob, log); err != nil {
			log.WithError(err).Error("Failed to update report state on prowjob")
			return nil, err
		}
	}

	return nil, nil
}

func (r *reconciler) shouldHandle(pj *prowv1.ProwJob) bool {
	refs := pj.Spec.ExtraRefs
	if pj.Spec.Refs != nil {
		refs = append(refs, *pj.Spec.Refs)
	}
	if len(refs) == 0 {
		return true
	}

	// It is possible to have conflicting settings here, we choose
	// to report if in doubt because reporting multiple times is
	// better than not reporting at all.
	var enabled bool
	for _, ref := range refs {
		if r.enablementChecker(ref.Org, ref.Repo) {
			enabled = true
			break
		}
	}

	return enabled
}
<commit_msg>Allow Crier to timeout from fetching/updating PJ.<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

// Package crier reports finished prowjob status to git providers.
package crier

import (
	"context"
	"fmt"
	"time"

	"github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/util/retry"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	ctrlruntimeclient "sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/manager"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"

	prowv1 "k8s.io/test-infra/prow/apis/prowjobs/v1"
)

type ReportClient interface {
	// Report reports a Prowjob. The provided logger is already populated with the
	// prowjob name and the reporter name.
	// If a reporter wants to defer reporting, it can return a reconcile.Result with a RequeueAfter
	Report(log *logrus.Entry, pj *prowv1.ProwJob) ([]*prowv1.ProwJob, *reconcile.Result, error)
	GetName() string
	// ShouldReport determines if a ProwJob should be reported. The provided logger
	// is already populated with the prowjob name and the reporter name.
	ShouldReport(log *logrus.Entry, pj *prowv1.ProwJob) bool
}

// reconciler struct defines how a controller should encapsulate
// logging, client connectivity, informing (list and watching)
// queueing, and handling of resource changes
type reconciler struct {
	pjclientset       ctrlruntimeclient.Client
	reporter          ReportClient
	enablementChecker func(org, repo string) bool
}

// New constructs a new instance of the crier reconciler.
func New(
	mgr manager.Manager,
	reporter ReportClient,
	numWorkers int,
	enablementChecker func(org, repo string) bool,
) error {
	if err := builder.
		ControllerManagedBy(mgr).
		// Is used for metrics, hence must be unique per controller instance
		Named(fmt.Sprintf("crier_%s", reporter.GetName())).
		For(&prowv1.ProwJob{}).
		WithOptions(controller.Options{MaxConcurrentReconciles: numWorkers}).
		Complete(&reconciler{
			pjclientset:       mgr.GetClient(),
			reporter:          reporter,
			enablementChecker: enablementChecker,
		}); err != nil {
		return fmt.Errorf("failed to construct controller: %w", err)
	}

	return nil
}

func (r *reconciler) updateReportState(ctx context.Context, pj *prowv1.ProwJob, log *logrus.Entry, reportedState prowv1.ProwJobState) error {
	// update pj report status
	newpj := pj.DeepCopy()
	// we set omitempty on PrevReportStates, so here we need to init it if is nil
	if newpj.Status.PrevReportStates == nil {
		newpj.Status.PrevReportStates = map[string]prowv1.ProwJobState{}
	}
	newpj.Status.PrevReportStates[r.reporter.GetName()] = reportedState

	if err := r.pjclientset.Patch(ctx, newpj, ctrlruntimeclient.MergeFrom(pj)); err != nil {
		return fmt.Errorf("failed to patch: %w", err)
	}

	// Block until the update is in the lister to make sure that events from another controller
	// that also does reporting dont trigger another report because our lister doesn't yet contain
	// the updated Status
	name := types.NamespacedName{Namespace: pj.Namespace, Name: pj.Name}
	if err := wait.Poll(100*time.Millisecond, 3*time.Second, func() (bool, error) {
		if err := r.pjclientset.Get(ctx, name, pj); err != nil {
			return false, err
		}
		if pj.Status.PrevReportStates != nil &&
			pj.Status.PrevReportStates[r.reporter.GetName()] == reportedState {
			return true, nil
		}
		return false, nil
	}); err != nil {
		return fmt.Errorf("failed to wait for updated report status to be in lister: %w", err)
	}
	return nil
}

func (r *reconciler) updateReportStateWithRetries(ctx context.Context, pj *prowv1.ProwJob, log *logrus.Entry) error {
	reportState := pj.Status.State
	log = log.WithFields(logrus.Fields{
		"prowjob":   pj.Name,
		"jobName":   pj.Spec.Job,
		"jobStatus": reportState,
	})
	// We have to retry here, if we return we lose the information that we already reported this job.
	if err := retry.RetryOnConflict(retry.DefaultBackoff, func() error {
		// Get it first, this is very cheap
		name := types.NamespacedName{Namespace: pj.Namespace, Name: pj.Name}
		if err := r.pjclientset.Get(ctx, name, pj); err != nil {
			return err
		}
		// Must not wrap until we have kube 1.19, otherwise the RetryOnConflict won't recognize conflicts
		// correctly
		return r.updateReportState(ctx, pj, log, reportState)
	}); err != nil {
		// Very subpar, we will report again. But even if we didn't do that now, we would do so
		// latest when crier gets restarted. In an ideal world, all reporters are idempotent and
		// reporting has no cost.
		return fmt.Errorf("failed to update report state on prowjob: %w", err)
	}

	log.Info("Successfully updated report state on prowjob")
	return nil
}

// Reconcile retrieves each queued item and takes the necessary handler action based off of if
// the item was created or deleted.
func (r *reconciler) Reconcile(ctx context.Context, req reconcile.Request) (reconcile.Result, error) {
	log := logrus.WithField("reporter", r.reporter.GetName()).WithField("key", req.String()).WithField("prowjob", req.Name)
	log.Debug("processing next key")
	// Limit reconciliation time to 15 minutes. This should more than enough time
	// for any reasonable reporter. Most reporters should set a stricter timeout
	// themselves. This mainly helps avoid leaking reconciliation threads that
	// will never complete.
	ctx, cancel := context.WithTimeout(ctx, 15*time.Minute)
	defer cancel()
	result, err := r.reconcile(ctx, log, req)
	if err != nil {
		log.WithError(err).Error("Reconciliation failed")
	}
	if result == nil {
		result = &reconcile.Result{}
	}
	return *result, err
}

func (r *reconciler) reconcile(ctx context.Context, log *logrus.Entry, req reconcile.Request) (*reconcile.Result, error) {

	pjContext, pjCancel := context.WithTimeout(ctx, 30*time.Second)
	defer pjCancel()
	var pj prowv1.ProwJob
	if err := r.pjclientset.Get(pjContext, req.NamespacedName, &pj); err != nil {
		if errors.IsNotFound(err) {
			log.Debug("object no longer exist")
			return nil, nil
		}

		return nil, fmt.Errorf("failed to get prowjob %s: %w", req.String(), err)
	}
	pjCancel()

	if !r.shouldHandle(&pj) {
		return nil, nil
	}

	log = log.WithField("jobName", pj.Spec.Job)

	if !pj.Spec.Report || !r.reporter.ShouldReport(log, &pj) {
		return nil, nil
	}

	// we set omitempty on PrevReportStates, so here we need to init it if is nil
	if pj.Status.PrevReportStates == nil {
		pj.Status.PrevReportStates = map[string]prowv1.ProwJobState{}
	}

	// already reported current state
	if pj.Status.PrevReportStates[r.reporter.GetName()] == pj.Status.State {
		log.Trace("Already reported")
		return nil, nil
	}

	log = log.WithField("jobStatus", pj.Status.State)
	log.Info("Will report state")
	pjs, requeue, err := r.reporter.Report(log, &pj)
	if err != nil {
		log.WithError(err).Error("failed to report job")
		return nil, fmt.Errorf("failed to report job: %w", err)
	}
	if requeue != nil {
		return requeue, nil
	}

	log.Infof("Reported job(s), now will update %d pj(s)", len(pjs))
	// Spend up to 10 mins attempting to update report state since the behavior
	// is poor if we fail. This long of a delay is needed for when the API
	// server is temporarily unavailable.
	updateContext, updateCancel := context.WithTimeout(ctx, 10*time.Minute)
	defer updateCancel()
	for _, pjob := range pjs {
		if err := r.updateReportStateWithRetries(updateContext, pjob, log); err != nil {
			log.WithError(err).Error("Failed to update report state on prowjob")
			return nil, err
		}
	}

	return nil, nil
}

func (r *reconciler) shouldHandle(pj *prowv1.ProwJob) bool {
	refs := pj.Spec.ExtraRefs
	if pj.Spec.Refs != nil {
		refs = append(refs, *pj.Spec.Refs)
	}
	if len(refs) == 0 {
		return true
	}

	// It is possible to have conflicting settings here, we choose
	// to report if in doubt because reporting multiple times is
	// better than not reporting at all.
	var enabled bool
	for _, ref := range refs {
		if r.enablementChecker(ref.Org, ref.Repo) {
			enabled = true
			break
		}
	}

	return enabled
}
<|endoftext|>


<commit_before>package freedom

import (
	"net"

	"github.com/v2ray/v2ray-core"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
)

type FreedomConnection struct {
	packet v2net.Packet
}

func NewFreedomConnection(firstPacket v2net.Packet) *FreedomConnection {
	return &FreedomConnection{
		packet: firstPacket,
	}
}

func (vconn *FreedomConnection) Start(ray core.OutboundRay) error {
	conn, err := net.Dial(vconn.packet.Destination().Network(), vconn.packet.Destination().Address().String())
	log.Info("Freedom: Opening connection to %s", vconn.packet.Destination().String())
	if err != nil {
		if ray != nil {
			close(ray.OutboundOutput())
		}
		return log.Error("Freedom: Failed to open connection: %s : %v", vconn.packet.Destination().String(), err)
	}

	if chunk := vconn.packet.Chunk(); chunk != nil {
		conn.Write(chunk)
	}

	if !vconn.packet.MoreChunks() {
		if ray != nil {
			close(ray.OutboundOutput())
		}
		return nil
	}

	input := ray.OutboundInput()
	output := ray.OutboundOutput()
	readFinish := make(chan bool)
	writeFinish := make(chan bool)

	go dumpInput(conn, input, writeFinish)
	go dumpOutput(conn, output, readFinish)

  go func() {
    <-writeFinish
    if tcpConn, ok := conn.(*net.TCPConn); ok {
      tcpConn.CloseWrite()
    }
    <-readFinish
    conn.Close()
  }()

	return nil
}

func dumpInput(conn net.Conn, input <-chan []byte, finish chan<- bool) {
	v2net.ChanToWriter(conn, input)
	close(finish)
}

func dumpOutput(conn net.Conn, output chan<- []byte, finish chan<- bool) {
	v2net.ReaderToChan(output, conn)
	close(output)
	close(finish)
}
<commit_msg>format code<commit_after>package freedom

import (
	"net"

	"github.com/v2ray/v2ray-core"
	"github.com/v2ray/v2ray-core/common/log"
	v2net "github.com/v2ray/v2ray-core/common/net"
)

type FreedomConnection struct {
	packet v2net.Packet
}

func NewFreedomConnection(firstPacket v2net.Packet) *FreedomConnection {
	return &FreedomConnection{
		packet: firstPacket,
	}
}

func (vconn *FreedomConnection) Start(ray core.OutboundRay) error {
	conn, err := net.Dial(vconn.packet.Destination().Network(), vconn.packet.Destination().Address().String())
	log.Info("Freedom: Opening connection to %s", vconn.packet.Destination().String())
	if err != nil {
		if ray != nil {
			close(ray.OutboundOutput())
		}
		return log.Error("Freedom: Failed to open connection: %s : %v", vconn.packet.Destination().String(), err)
	}

	if chunk := vconn.packet.Chunk(); chunk != nil {
		conn.Write(chunk)
	}

	if !vconn.packet.MoreChunks() {
		if ray != nil {
			close(ray.OutboundOutput())
		}
		return nil
	}

	input := ray.OutboundInput()
	output := ray.OutboundOutput()
	readFinish := make(chan bool)
	writeFinish := make(chan bool)

	go dumpInput(conn, input, writeFinish)
	go dumpOutput(conn, output, readFinish)

	go func() {
		<-writeFinish
		if tcpConn, ok := conn.(*net.TCPConn); ok {
			tcpConn.CloseWrite()
		}
		<-readFinish
		conn.Close()
	}()

	return nil
}

func dumpInput(conn net.Conn, input <-chan []byte, finish chan<- bool) {
	v2net.ChanToWriter(conn, input)
	close(finish)
}

func dumpOutput(conn net.Conn, output chan<- []byte, finish chan<- bool) {
	v2net.ReaderToChan(output, conn)
	close(output)
	close(finish)
}
<|endoftext|>


<commit_before><commit_msg>tcp connection shutdown and close<commit_after><|endoftext|>