<commit_before>// $ go run gazel.go -r ./vendor
// $ go run gazel.go -r ./federation
// $ go run gazel.go -r ./contrib
// $ go run gazel.go -r ./plugin
// $ go run gazel.go -r ./test
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"sort"
	"strings"

	bzl "github.com/bazelbuild/buildifier/core"
	"github.com/golang/glog"
)

var kubeRoot = os.Getenv("KUBE_ROOT")
var dryRun = flag.Bool("dry-run", false, "run in dry mode")

func main() {
	flag.Parse()
	flag.Set("alsologtostderr", "true")
	v := Venderor{
		ctx: &build.Default,
	}
	if len(os.Args) == 2 {
		v.updateSinglePkg(os.Args[1])
	} else {
		v.walkVendor()
		if err := v.walkRepo(); err != nil {
			glog.Fatalf("err: %v", err)
		}
	}
}

type Venderor struct {
	ctx *build.Context
}

func writeHeaders(file *bzl.File) {
	pkgRule := bzl.Rule{
		&bzl.CallExpr{
			X: &bzl.LiteralExpr{Token: "package"},
		},
	}
	pkgRule.SetAttr("default_visibility", asExpr([]string{"//visibility:public"}))

	file.Stmt = append(file.Stmt,
		[]bzl.Expr{
			pkgRule.Call,
			&bzl.CallExpr{
				X:    &bzl.LiteralExpr{Token: "licenses"},
				List: []bzl.Expr{asExpr([]string{"notice"})},
			},
			&bzl.CallExpr{
				X: &bzl.LiteralExpr{Token: "load"},
				List: asExpr([]string{
					"@io_bazel_rules_go//go:def.bzl",
					"go_binary",
					"go_library",
					"go_test",
					"cgo_library",
				}).(*bzl.ListExpr).List,
			},
		}...,
	)
}

func writeRules(file *bzl.File, rules []*bzl.Rule) {
	for _, rule := range rules {
		file.Stmt = append(file.Stmt, rule.Call)
	}
}

func (v *Venderor) resolve(ipath string) Label {
	if strings.HasPrefix(ipath, "k8s.io/kubernetes") {
		return Label{
			pkg: strings.TrimPrefix(ipath, "k8s.io/kubernetes/"),
			tag: "go_default_library",
		}
	}
	return Label{
		pkg: "vendor",
		tag: ipath,
	}
}

func (v *Venderor) walk(root string, f func(path, ipath string, pkg *build.Package) error) error {
	return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		ipath, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		if !info.IsDir() {
			return nil
		}
		pkg, err := v.ctx.ImportDir(filepath.Join(kubeRoot, path), build.ImportComment)
		if err != nil {
			if _, ok := err.(*build.NoGoError); err != nil && ok {
				return nil
			} else {
				return err
			}
		}

		return f(path, ipath, pkg)
	})
}

func (v *Venderor) walkRepo() error {
	for _, root := range []string{
		"./pkg",
		"./cmd",
		"./third_party",
		"./plugin",
		"./test",
		"./federation",
	} {
		if err := v.walk(root, v.updatePkg); err != nil {
			return err
		}
	}
	return nil
}

func (v *Venderor) updateSinglePkg(path string) error {
	pkg, err := v.ctx.ImportDir("./"+path, build.ImportComment)
	if err != nil {
		if _, ok := err.(*build.NoGoError); err != nil && ok {
			return nil
		} else {
			return err
		}
	}
	return v.updatePkg(path, "", pkg)
}

func (v *Venderor) updatePkg(path, _ string, pkg *build.Package) error {
	var rules []*bzl.Rule

	var attrs Attrs = make(Attrs)
	srcs := asExpr(merge(pkg.GoFiles, pkg.SFiles)).(*bzl.ListExpr)

	deps := v.extractDeps(pkg.Imports)

	if len(srcs.List) == 0 {
		return nil
	}
	attrs.Set("srcs", srcs)

	if len(deps.List) > 0 {
		attrs.Set("deps", deps)
	}

	if pkg.IsCommand() {
		rules = append(rules, newRule("go_binary", filepath.Base(pkg.Dir), attrs))
	} else {
		rules = append(rules, newRule("go_library", "go_default_library", attrs))
		if len(pkg.TestGoFiles) != 0 {
			rules = append(rules, newRule("go_test", "go_default_test", map[string]bzl.Expr{
				"srcs":    asExpr(pkg.TestGoFiles),
				"deps":    v.extractDeps(pkg.TestImports),
				"library": asExpr("go_default_library"),
			}))
		}
	}

	if len(pkg.XTestGoFiles) != 0 {
		rules = append(rules, newRule("go_test", "go_default_xtest", map[string]bzl.Expr{
			"srcs": asExpr(pkg.XTestGoFiles),
			"deps": v.extractDeps(pkg.XTestImports),
		}))
	}

	wrote, err := ReconcileRules(filepath.Join(path, "BUILD"), rules)
	if err != nil {
		return err
	}
	if wrote {
		fmt.Fprintf(os.Stderr, "wrote BUILD for %q\n", pkg.ImportPath)
	}
	return nil
}

func (v *Venderor) walkVendor() {
	var rules []*bzl.Rule
	if err := v.walk("./vendor", func(path, ipath string, pkg *build.Package) error {
		var attrs Attrs = make(Attrs)

		srcs := asExpr(
			apply(
				merge(pkg.GoFiles, pkg.SFiles),
				mapper(func(s string) string {
					return strings.TrimPrefix(filepath.Join(path, s), "vendor/")
				}),
			),
		).(*bzl.ListExpr)

		cgoSrcs := asExpr(
			apply(
				merge(pkg.CgoFiles, pkg.CFiles, pkg.CXXFiles, pkg.HFiles),
				mapper(func(s string) string {
					return strings.TrimPrefix(filepath.Join(path, s), "vendor/")
				}),
			),
		).(*bzl.ListExpr)

		deps := v.extractDeps(pkg.Imports)
		attrs.Set("srcs", srcs)

		if len(deps.List) > 0 {
			attrs.Set("deps", deps)
		}

		if pkg.IsCommand() {
			rules = append(rules, newRule("go_binary", v.resolve(ipath).tag, attrs))
		} else {
			if len(cgoSrcs.List) != 0 {
				cgoPname := v.resolve(ipath).tag + "_cgo"
				cgoDeps := v.extractDeps(pkg.TestImports)
				cgoRule := newRule("cgo_library", cgoPname, map[string]bzl.Expr{
					"srcs":      cgoSrcs,
					"clinkopts": asExpr([]string{"-ldl", "-lz", "-lm", "-lpthread", "-ldl"}),
				})
				rules = append(rules, cgoRule)
				if len(cgoDeps.List) != 0 {
					cgoRule.SetAttr("deps", cgoDeps)
				}
				attrs["library"] = asExpr(cgoPname)
			}
			rules = append(rules, newRule("go_library", v.resolve(ipath).tag, attrs))
		}
		return nil
	}); err != nil {
		glog.Fatalf("err: %v", err)
	}
	if _, err := ReconcileRules("./vendor/BUILD", rules); err != nil {
		glog.Fatalf("err: %v", err)
	}
}

func (v *Venderor) extractDeps(deps []string) *bzl.ListExpr {
	return asExpr(
		apply(
			merge(deps),
			filterer(func(s string) bool {
				pkg, err := v.ctx.Import(s, kubeRoot, build.ImportComment)
				if err != nil {
					return false
				}
				if pkg.Goroot {
					return false
				}
				return true
			}),
			mapper(func(s string) string {
				return v.resolve(s).String()
			}),
		),
	).(*bzl.ListExpr)
}

type Attrs map[string]bzl.Expr

func (a Attrs) Set(name string, expr bzl.Expr) {
	a[name] = expr
}

type Label struct {
	pkg, tag string
}

func (l Label) String() string {
	return fmt.Sprintf("//%v:%v", l.pkg, l.tag)
}

func asExpr(e interface{}) bzl.Expr {
	rv := reflect.ValueOf(e)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return &bzl.LiteralExpr{Token: fmt.Sprintf("%d", e)}
	case reflect.Float32, reflect.Float64:
		return &bzl.LiteralExpr{Token: fmt.Sprintf("%f", e)}
	case reflect.String:
		return &bzl.StringExpr{Value: e.(string)}
	case reflect.Slice, reflect.Array:
		var list []bzl.Expr
		for i := 0; i < rv.Len(); i++ {
			list = append(list, asExpr(rv.Index(i).Interface()))
		}
		return &bzl.ListExpr{List: list}
	default:
		glog.Fatalf("Uh oh")
		return nil
	}
}

type Sed func(s []string) []string

func mapString(in []string, f func(string) string) []string {
	var out []string
	for _, s := range in {
		out = append(out, f(s))
	}
	return out
}

func mapper(f func(string) string) Sed {
	return func(in []string) []string {
		return mapString(in, f)
	}
}

func filterString(in []string, f func(string) bool) []string {
	var out []string
	for _, s := range in {
		if f(s) {
			out = append(out, s)
		}
	}
	return out
}

func filterer(f func(string) bool) Sed {
	return func(in []string) []string {
		return filterString(in, f)
	}
}

func apply(stream []string, seds ...Sed) []string {
	for _, sed := range seds {
		stream = sed(stream)
	}
	return stream
}

func merge(streams ...[]string) []string {
	var out []string
	for _, stream := range streams {
		out = append(out, stream...)
	}
	return out
}

func newRule(kind, name string, attrs map[string]bzl.Expr) *bzl.Rule {
	rule := &bzl.Rule{
		Call: &bzl.CallExpr{
			X: &bzl.LiteralExpr{Token: kind},
		},
	}
	rule.SetAttr("name", asExpr(name))
	for _, k := range sortKeys(attrs) {
		rule.SetAttr(k, attrs[k])
	}
	rule.SetAttr("tags", asExpr([]string{"automanaged"}))
	return rule
}

func sortKeys(attrs map[string]bzl.Expr) []string {
	ks := []string{}
	for k, _ := range attrs {
		ks = append(ks, k)
	}
	sort.Sort(orderedAttrs(ks))
	return ks
}

type orderedAttrs []string

var orderedAttrsPriority = map[string]int{
	"name":    1,
	"srcs":    2,
	"deps":    3,
	"library": 4,
	"tags":    5,
}

func (o orderedAttrs) Len() int {
	return len(o)
}

func (o orderedAttrs) Swap(i, j int) {
	o[i], o[j] = o[j], o[i]
}

func (o orderedAttrs) Less(i, j int) bool {
	ip, iok := orderedAttrsPriority[o[i]]
	jp, jok := orderedAttrsPriority[o[j]]
	if iok && jok {
		return ip < jp
	}
	if iok && !jok {
		return false
	}
	if jok && !iok {
		return true
	}
	return o[i] < o[j]
}

func ReconcileRules(path string, rules []*bzl.Rule) (bool, error) {
	info, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		f := &bzl.File{}
		writeHeaders(f)
		writeRules(f, rules)
		return true, writeFile(path, bzl.Format(f))
	} else if err != nil {
		return false, err
	}
	if info.IsDir() {
		return false, fmt.Errorf("%q cannot be a directory", path)
	}
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return false, err
	}
	f, err := bzl.Parse(path, b)
	if err != nil {
		return false, err
	}
	oldRules := make(map[string]*bzl.Rule)
	for _, r := range f.Rules("") {
		oldRules[r.Name()] = r
	}
	for _, r := range rules {
		o, ok := oldRules[r.Name()]
		if !ok {
			f.Stmt = append(f.Stmt, r.Call)
			continue
		}
		if !RuleIsManaged(o) {
			continue
		}
		reconcileAttr := func(o, n *bzl.Rule, name string) {
			if e := n.Attr(name); e != nil {
				o.SetAttr(name, e)
			} else {
				o.DelAttr(name)
			}
		}
		reconcileAttr(o, r, "srcs")
		reconcileAttr(o, r, "deps")
		reconcileAttr(o, r, "library")
		delete(oldRules, r.Name())
	}
	for _, r := range oldRules {
		if !RuleIsManaged(r) {
			continue
		}
		f.DelRules(r.Kind(), r.Name())
	}
	out := bzl.Format(f)
	orig, err := ioutil.ReadFile(path)
	if err != nil {
		return false, err
	}
	if bytes.Compare(out, orig) == 0 {
		return false, err
	}
	return true, writeFile(path, bzl.Format(f))
}

func RuleIsManaged(r *bzl.Rule) bool {
	var automanaged bool
	for _, tag := range r.AttrStrings("tags") {
		if tag == "automanaged" {
			automanaged = true
			break
		}
	}
	return automanaged
}

func writeFile(path string, b []byte) error {
	if *dryRun {
		return nil
	}
	return ioutil.WriteFile(path, b, 0644)

}
<commit_msg>fix flag bug<commit_after>// $ go run gazel.go -r ./vendor
// $ go run gazel.go -r ./federation
// $ go run gazel.go -r ./contrib
// $ go run gazel.go -r ./plugin
// $ go run gazel.go -r ./test
package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"reflect"
	"strings"

	bzl "github.com/bazelbuild/buildifier/core"
	"github.com/golang/glog"
)

var kubeRoot = os.Getenv("KUBE_ROOT")
var dryRun = flag.Bool("dry-run", false, "run in dry mode")

func main() {
	flag.Parse()
	flag.Set("alsologtostderr", "true")
	v := Venderor{
		ctx: &build.Default,
	}
	if len(flag.Args()) == 2 {
		v.updateSinglePkg(flag.Args()[1])
	} else {
		v.walkVendor()
		if err := v.walkRepo(); err != nil {
			glog.Fatalf("err: %v", err)
		}
	}
}

type Venderor struct {
	ctx *build.Context
}

func writeHeaders(file *bzl.File) {
	pkgRule := bzl.Rule{
		&bzl.CallExpr{
			X: &bzl.LiteralExpr{Token: "package"},
		},
	}
	pkgRule.SetAttr("default_visibility", asExpr([]string{"//visibility:public"}))

	file.Stmt = append(file.Stmt,
		[]bzl.Expr{
			pkgRule.Call,
			&bzl.CallExpr{
				X:    &bzl.LiteralExpr{Token: "licenses"},
				List: []bzl.Expr{asExpr([]string{"notice"})},
			},
			&bzl.CallExpr{
				X: &bzl.LiteralExpr{Token: "load"},
				List: asExpr([]string{
					"@io_bazel_rules_go//go:def.bzl",
					"go_binary",
					"go_library",
					"go_test",
					"cgo_library",
				}).(*bzl.ListExpr).List,
			},
		}...,
	)
}

func writeRules(file *bzl.File, rules []*bzl.Rule) {
	for _, rule := range rules {
		file.Stmt = append(file.Stmt, rule.Call)
	}
}

func (v *Venderor) resolve(ipath string) Label {
	if strings.HasPrefix(ipath, "k8s.io/kubernetes") {
		return Label{
			pkg: strings.TrimPrefix(ipath, "k8s.io/kubernetes/"),
			tag: "go_default_library",
		}
	}
	return Label{
		pkg: "vendor",
		tag: ipath,
	}
}

func (v *Venderor) walk(root string, f func(path, ipath string, pkg *build.Package) error) error {
	return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		ipath, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		if !info.IsDir() {
			return nil
		}
		pkg, err := v.ctx.ImportDir(filepath.Join(kubeRoot, path), build.ImportComment)
		if err != nil {
			if _, ok := err.(*build.NoGoError); err != nil && ok {
				return nil
			} else {
				return err
			}
		}

		return f(path, ipath, pkg)
	})
}

func (v *Venderor) walkRepo() error {
	for _, root := range []string{
		"./pkg",
		"./cmd",
		"./third_party",
		"./plugin",
		"./test",
		"./federation",
	} {
		if err := v.walk(root, v.updatePkg); err != nil {
			return err
		}
	}
	return nil
}

func (v *Venderor) updateSinglePkg(path string) error {
	pkg, err := v.ctx.ImportDir("./"+path, build.ImportComment)
	if err != nil {
		if _, ok := err.(*build.NoGoError); err != nil && ok {
			return nil
		} else {
			return err
		}
	}
	return v.updatePkg(path, "", pkg)
}

func (v *Venderor) updatePkg(path, _ string, pkg *build.Package) error {
	var rules []*bzl.Rule

	var attrs Attrs = make(Attrs)
	srcs := asExpr(merge(pkg.GoFiles, pkg.SFiles)).(*bzl.ListExpr)

	deps := v.extractDeps(pkg.Imports)

	if len(srcs.List) == 0 {
		return nil
	}
	attrs.Set("srcs", srcs)

	if len(deps.List) > 0 {
		attrs.Set("deps", deps)
	}

	if pkg.IsCommand() {
		rules = append(rules, newRule("go_binary", filepath.Base(pkg.Dir), attrs))
	} else {
		rules = append(rules, newRule("go_library", "go_default_library", attrs))
		if len(pkg.TestGoFiles) != 0 {
			rules = append(rules, newRule("go_test", "go_default_test", map[string]bzl.Expr{
				"srcs":    asExpr(pkg.TestGoFiles),
				"deps":    v.extractDeps(pkg.TestImports),
				"library": asExpr("go_default_library"),
			}))
		}
	}

	if len(pkg.XTestGoFiles) != 0 {
		rules = append(rules, newRule("go_test", "go_default_xtest", map[string]bzl.Expr{
			"srcs": asExpr(pkg.XTestGoFiles),
			"deps": v.extractDeps(pkg.XTestImports),
		}))
	}

	wrote, err := ReconcileRules(filepath.Join(path, "BUILD"), rules)
	if err != nil {
		return err
	}
	if wrote {
		fmt.Fprintf(os.Stderr, "wrote BUILD for %q\n", pkg.ImportPath)
	}
	return nil
}

func (v *Venderor) walkVendor() {
	var rules []*bzl.Rule
	if err := v.walk("./vendor", func(path, ipath string, pkg *build.Package) error {
		var attrs Attrs = make(Attrs)

		srcs := asExpr(
			apply(
				merge(pkg.GoFiles, pkg.SFiles),
				mapper(func(s string) string {
					return strings.TrimPrefix(filepath.Join(path, s), "vendor/")
				}),
			),
		).(*bzl.ListExpr)

		cgoSrcs := asExpr(
			apply(
				merge(pkg.CgoFiles, pkg.CFiles, pkg.CXXFiles, pkg.HFiles),
				mapper(func(s string) string {
					return strings.TrimPrefix(filepath.Join(path, s), "vendor/")
				}),
			),
		).(*bzl.ListExpr)

		deps := v.extractDeps(pkg.Imports)
		attrs.Set("srcs", srcs)

		if len(deps.List) > 0 {
			attrs.Set("deps", deps)
		}

		if pkg.IsCommand() {
			rules = append(rules, newRule("go_binary", v.resolve(ipath).tag, attrs))
		} else {
			if len(cgoSrcs.List) != 0 {
				cgoPname := v.resolve(ipath).tag + "_cgo"
				cgoDeps := v.extractDeps(pkg.TestImports)
				cgoRule := newRule("cgo_library", cgoPname, map[string]bzl.Expr{
					"srcs":      cgoSrcs,
					"clinkopts": asExpr([]string{"-ldl", "-lz", "-lm", "-lpthread", "-ldl"}),
				})
				rules = append(rules, cgoRule)
				if len(cgoDeps.List) != 0 {
					cgoRule.SetAttr("deps", cgoDeps)
				}
				attrs["library"] = asExpr(cgoPname)
			}
			rules = append(rules, newRule("go_library", v.resolve(ipath).tag, attrs))
		}
		return nil
	}); err != nil {
		glog.Fatalf("err: %v", err)
	}
	if _, err := ReconcileRules("./vendor/BUILD", rules); err != nil {
		glog.Fatalf("err: %v", err)
	}
}

func (v *Venderor) extractDeps(deps []string) *bzl.ListExpr {
	return asExpr(
		apply(
			merge(deps),
			filterer(func(s string) bool {
				pkg, err := v.ctx.Import(s, kubeRoot, build.ImportComment)
				if err != nil {
					return false
				}
				if pkg.Goroot {
					return false
				}
				return true
			}),
			mapper(func(s string) string {
				return v.resolve(s).String()
			}),
		),
	).(*bzl.ListExpr)
}

type Attrs map[string]bzl.Expr

func (a Attrs) Set(name string, expr bzl.Expr) {
	a[name] = expr
}

type Label struct {
	pkg, tag string
}

func (l Label) String() string {
	return fmt.Sprintf("//%v:%v", l.pkg, l.tag)
}

func asExpr(e interface{}) bzl.Expr {
	rv := reflect.ValueOf(e)
	switch rv.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
		reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return &bzl.LiteralExpr{Token: fmt.Sprintf("%d", e)}
	case reflect.Float32, reflect.Float64:
		return &bzl.LiteralExpr{Token: fmt.Sprintf("%f", e)}
	case reflect.String:
		return &bzl.StringExpr{Value: e.(string)}
	case reflect.Slice, reflect.Array:
		var list []bzl.Expr
		for i := 0; i < rv.Len(); i++ {
			list = append(list, asExpr(rv.Index(i).Interface()))
		}
		return &bzl.ListExpr{List: list}
	default:
		glog.Fatalf("Uh oh")
		return nil
	}
}

type Sed func(s []string) []string

func mapString(in []string, f func(string) string) []string {
	var out []string
	for _, s := range in {
		out = append(out, f(s))
	}
	return out
}

func mapper(f func(string) string) Sed {
	return func(in []string) []string {
		return mapString(in, f)
	}
}

func filterString(in []string, f func(string) bool) []string {
	var out []string
	for _, s := range in {
		if f(s) {
			out = append(out, s)
		}
	}
	return out
}

func filterer(f func(string) bool) Sed {
	return func(in []string) []string {
		return filterString(in, f)
	}
}

func apply(stream []string, seds ...Sed) []string {
	for _, sed := range seds {
		stream = sed(stream)
	}
	return stream
}

func merge(streams ...[]string) []string {
	var out []string
	for _, stream := range streams {
		out = append(out, stream...)
	}
	return out
}

func newRule(kind, name string, attrs map[string]bzl.Expr) *bzl.Rule {
	rule := &bzl.Rule{
		Call: &bzl.CallExpr{
			X: &bzl.LiteralExpr{Token: kind},
		},
	}
	rule.SetAttr("name", asExpr(name))
	for k, v := range attrs {
		rule.SetAttr(k, v)
	}
	rule.SetAttr("tags", asExpr([]string{"automanaged"}))
	return rule
}

func ReconcileRules(path string, rules []*bzl.Rule) (bool, error) {
	info, err := os.Stat(path)
	if err != nil && os.IsNotExist(err) {
		f := &bzl.File{}
		writeHeaders(f)
		writeRules(f, rules)
		return writeFile(path, f, false)
	} else if err != nil {
		return false, err
	}
	if info.IsDir() {
		return false, fmt.Errorf("%q cannot be a directory", path)
	}
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return false, err
	}
	f, err := bzl.Parse(path, b)
	if err != nil {
		return false, err
	}
	oldRules := make(map[string]*bzl.Rule)
	for _, r := range f.Rules("") {
		oldRules[r.Name()] = r
	}
	for _, r := range rules {
		o, ok := oldRules[r.Name()]
		if !ok {
			f.Stmt = append(f.Stmt, r.Call)
			continue
		}
		if !RuleIsManaged(o) {
			continue
		}
		reconcileAttr := func(o, n *bzl.Rule, name string) {
			if e := n.Attr(name); e != nil {
				o.SetAttr(name, e)
			} else {
				o.DelAttr(name)
			}
		}
		reconcileAttr(o, r, "srcs")
		reconcileAttr(o, r, "deps")
		reconcileAttr(o, r, "library")
		delete(oldRules, r.Name())
	}
	for _, r := range oldRules {
		if !RuleIsManaged(r) {
			continue
		}
		f.DelRules(r.Kind(), r.Name())
	}
	return writeFile(path, f, true)
}

func RuleIsManaged(r *bzl.Rule) bool {
	var automanaged bool
	for _, tag := range r.AttrStrings("tags") {
		if tag == "automanaged" {
			automanaged = true
			break
		}
	}
	return automanaged
}

func writeFile(path string, f *bzl.File, exists bool) (bool, error) {
	var info bzl.RewriteInfo
	bzl.Rewrite(f, &info)
	out := bzl.Format(f)
	if exists {
		orig, err := ioutil.ReadFile(path)
		if err != nil {
			return false, err
		}
		if bytes.Compare(out, orig) == 0 {
			return false, nil
		}
	}
	if *dryRun {
		return true, nil
	}
	return true, ioutil.WriteFile(path, out, 0644)

}
<|endoftext|>


<commit_before>package kite

import (
	"errors"
	"fmt"
	"github.com/op/go-logging"
	"koding/newkite/dnode"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"strconv"
	"sync"
	"time"
)

const DefaultCallTimeout = 4000 // milliseconds

// RemoteKite is the client for communicating with another Kite.
// It has Call() and Go() methods for calling methods sync/async way.
type RemoteKite struct {
	// The information about the kite that we are connecting to.
	protocol.Kite

	// A reference to the current Kite running.
	localKite *Kite

	// A reference to the Kite's logger for easy access.
	Log *logging.Logger

	// Credentials that we sent in each request.
	Authentication callAuthentication

	// dnode RPC client that processes messages.
	client *rpc.Client

	// To signal waiters of Go() on disconnect.
	disconnect chan bool

	// Duration to wait reply from remote when making a request with Call().
	callTimeout time.Duration
}

// NewRemoteKite returns a pointer to a new RemoteKite. The returned instance
// is not connected. You have to call Dial() or DialForever() before calling
// Call() and Go() methods.
func (k *Kite) NewRemoteKite(kite protocol.Kite, auth callAuthentication) *RemoteKite {
	r := &RemoteKite{
		Kite:           kite,
		localKite:      k,
		Log:            k.Log,
		Authentication: auth,
		client:         k.server.NewClientWithHandlers(),
		disconnect:     make(chan bool),
	}
	r.SetCallTimeout(DefaultCallTimeout)

	// We need a reference to the local kite when a method call is received.
	r.client.Properties()["localKite"] = k

	r.OnConnect(func() {
		if r.Authentication.ValidUntil == nil {
			return
		}

		// Start a goroutine that will renew the token before it expires.
		go r.tokenRenewer()
	})

	var m sync.Mutex
	r.OnDisconnect(func() {
		m.Lock()
		close(r.disconnect)
		r.disconnect = make(chan bool)
		m.Unlock()
	})

	return r
}

// newRemoteKiteWithClient returns a pointer to new RemoteKite instance.
// The client will be replaced with the given client.
// Used to give the Kite method handler a working RemoteKite to call methods
// on other side.
func (k *Kite) newRemoteKiteWithClient(kite protocol.Kite, auth callAuthentication, client *rpc.Client) *RemoteKite {
	r := k.NewRemoteKite(kite, auth)
	r.client = client
	r.client.Properties()["localKite"] = k
	return r
}

func (r *RemoteKite) SetCallTimeout(ms uint) {
	r.callTimeout = time.Duration(ms) * time.Millisecond
}

// Dial connects to the remote Kite. Returns error if it can't.
func (r *RemoteKite) Dial() (err error) {
	addr := r.Kite.Addr()
	r.Log.Info("Dialing remote kite: [%s %s]", r.Kite.Name, addr)
	return r.client.Dial("ws://" + addr + "/dnode")
}

// Dial connects to the remote Kite. If it can't connect, it retries indefinitely.
func (r *RemoteKite) DialForever() {
	addr := r.Kite.Addr()
	r.Log.Info("Dialing remote kite: [%s %s]", r.Kite.Name, addr)
	r.client.DialForever("ws://" + addr + "/dnode")
}

func (r *RemoteKite) Close() {
	r.client.Close()
}

// OnConnect registers a function to run on connect.
func (r *RemoteKite) OnConnect(handler func()) {
	r.client.OnConnect(handler)
}

// OnDisconnect registers a function to run on disconnect.
func (r *RemoteKite) OnDisconnect(handler func()) {
	r.client.OnDisconnect(handler)
}

func (r *RemoteKite) tokenRenewer() {
	for {
		renewTime := r.Authentication.ValidUntil.Add(-30 * time.Second)
		select {
		case <-time.After(renewTime.Sub(time.Now().UTC())):
			if err := r.renewTokenUntilDisconnect(); err != nil {
				return
			}
		case <-r.disconnect:
			return
		}
	}
}

// renewToken retries until the request is successful or disconnect.
func (r *RemoteKite) renewTokenUntilDisconnect() error {
	const retryInterval = 10 * time.Second

	if err := r.renewToken(); err == nil {
		return nil
	}

	for {
		select {
		case <-time.After(retryInterval):
			if err := r.renewToken(); err != nil {
				r.Log.Error("error: %s Cannot renew token for Kite: %s I will retry in %d seconds...", err.Error(), r.Kite.ID, retryInterval)
				continue
			}

			break
		case <-r.disconnect:
			return errors.New("disconnect")
		}
	}

	return nil
}

func (r *RemoteKite) renewToken() error {
	tkn, err := r.localKite.Kontrol.GetToken(&r.Kite)
	if err != nil {
		return err
	}

	validUntil := time.Now().UTC().Add(time.Duration(tkn.TTL) * time.Second)
	r.Authentication.Key = tkn.Key
	r.Authentication.ValidUntil = &validUntil

	return nil
}

// CallOptions is the type of first argument in the dnode message.
// Second argument is a callback function.
// It is used when unmarshalling a dnode message.
type CallOptions struct {
	// Arguments to the method
	WithArgs       *dnode.Partial     `json:"withArgs"`
	Kite           protocol.Kite      `json:"kite"`
	Authentication callAuthentication `json:"authentication"`
}

// callOptionsOut is the same structure with CallOptions.
// It is used when marshalling a dnode message.
type callOptionsOut struct {
	CallOptions
	// Override this when sending because args will not be a *dnode.Partial.
	WithArgs interface{} `json:"withArgs"`
}

// That's what we send as a first argument in dnode message.
func (r *RemoteKite) makeOptions(args interface{}) *callOptionsOut {
	return &callOptionsOut{
		WithArgs: args,
		CallOptions: CallOptions{
			Kite:           r.localKite.Kite,
			Authentication: r.Authentication,
		},
	}
}

type callAuthentication struct {
	// Type can be "kodingKey", "token" or "sessionID" for now.
	Type       string     `json:"type"`
	Key        string     `json:"key"`
	ValidUntil *time.Time `json:"-"`
}

type response struct {
	Result *dnode.Partial
	Err    error
}

// Call makes a blocking method call to the server.
// Waits until the callback function is called by the other side and
// returns the result and the error.
func (r *RemoteKite) Call(method string, args interface{}) (result *dnode.Partial, err error) {
	response := <-r.Go(method, args)
	return response.Result, response.Err
}

// Go makes an unblocking method call to the server.
// It returns a channel that the caller can wait on it to get the response.
func (r *RemoteKite) Go(method string, args interface{}) chan *response {
	// We will return this channel to the caller.
	// It can wait on this channel to get the response.
	r.Log.Debug("Calling method [%s] on kite [%s]", method, r.Name)
	responseChan := make(chan *response, 1)

	r.send(method, args, responseChan)

	return responseChan
}

// send sends the method with callback to the server.
func (r *RemoteKite) send(method string, args interface{}, responseChan chan *response) {
	// To clean the sent callback after response is received.
	// Send/Receive in a channel to prevent race condition because
	// the callback is run in a separate goroutine.
	removeCallback := make(chan uint64, 1)

	// When a callback is called it will send the response to this channel.
	doneChan := make(chan *response, 1)

	opts := r.makeOptions(args)
	cb := r.makeResponseCallback(doneChan, removeCallback)

	callbacks, err := r.client.Call(method, opts, cb)
	if err != nil {
		responseChan <- &response{
			Result: nil,
			Err: fmt.Errorf("Calling method [%s] on [%s] error: %s",
				method, r.Kite.Name, err),
		}
		return
	}

	// Waits until the response has came or the connection has disconnected.
	go func() {
		select {
		case <-r.disconnect:
			responseChan <- &response{nil, errors.New("Client disconnected")}
		case resp := <-doneChan:
			responseChan <- resp
		case <-time.After(r.callTimeout):
			responseChan <- &response{nil, errors.New("Timeout")}

			// Remove the callback function from the map so we do not
			// consume memory for unused callbacks.
			if id, ok := <-removeCallback; ok {
				r.client.RemoveCallback(id)
			}
		}
	}()

	sendCallbackID(callbacks, removeCallback)
}

// sendCallbackID send the callback number to be deleted after response is received.
func sendCallbackID(callbacks map[string]dnode.Path, ch chan uint64) {
	if len(callbacks) > 0 {
		// Find max callback ID.
		max := uint64(0)
		for id, _ := range callbacks {
			i, _ := strconv.ParseUint(id, 10, 64)
			if i > max {
				max = i
			}
		}

		ch <- max
	} else {
		close(ch)
	}
}

// makeResponseCallback prepares and returns a callback function sent to the server.
// The caller of the Call() is blocked until the server calls this callback function.
// Sets theResponse and notifies the caller by sending to done channel.
func (r *RemoteKite) makeResponseCallback(doneChan chan *response, removeCallback <-chan uint64) Callback {
	return Callback(func(request *Request) {
		var (
			err    error          // First argument
			result *dnode.Partial // Second argument
		)

		// Notify that the callback is finished.
		defer func() { doneChan <- &response{result, err} }()

		// Remove the callback function from the map so we do not
		// consume memory for unused callbacks.
		if id, ok := <-removeCallback; ok {
			r.client.RemoveCallback(id)
		}

		// Arguments to our response callback It is a slice of length 2.
		// The first argument is the error string,
		// the second argument is the result.
		responseArgs := request.Args.MustSliceOfLength(2)

		// The second argument is our result.
		result = responseArgs[1]

		// This is error argument. Unmarshal panics if it is null.
		if responseArgs[0] == nil {
			return
		}

		// Read the error argument in response.
		var kiteErr *Error
		err = responseArgs[0].Unmarshal(kiteErr)
		if err != nil {
			return
		}

		err = kiteErr
	})
}
<commit_msg>kite: add comment<commit_after>package kite

import (
	"errors"
	"fmt"
	"github.com/op/go-logging"
	"koding/newkite/dnode"
	"koding/newkite/dnode/rpc"
	"koding/newkite/protocol"
	"strconv"
	"sync"
	"time"
)

const DefaultCallTimeout = 4000 // milliseconds

// RemoteKite is the client for communicating with another Kite.
// It has Call() and Go() methods for calling methods sync/async way.
type RemoteKite struct {
	// The information about the kite that we are connecting to.
	protocol.Kite

	// A reference to the current Kite running.
	localKite *Kite

	// A reference to the Kite's logger for easy access.
	Log *logging.Logger

	// Credentials that we sent in each request.
	Authentication callAuthentication

	// dnode RPC client that processes messages.
	client *rpc.Client

	// To signal waiters of Go() on disconnect.
	disconnect chan bool

	// Duration to wait reply from remote when making a request with Call().
	callTimeout time.Duration
}

// NewRemoteKite returns a pointer to a new RemoteKite. The returned instance
// is not connected. You have to call Dial() or DialForever() before calling
// Call() and Go() methods.
func (k *Kite) NewRemoteKite(kite protocol.Kite, auth callAuthentication) *RemoteKite {
	r := &RemoteKite{
		Kite:           kite,
		localKite:      k,
		Log:            k.Log,
		Authentication: auth,
		client:         k.server.NewClientWithHandlers(),
		disconnect:     make(chan bool),
	}
	r.SetCallTimeout(DefaultCallTimeout)

	// We need a reference to the local kite when a method call is received.
	r.client.Properties()["localKite"] = k

	r.OnConnect(func() {
		if r.Authentication.ValidUntil == nil {
			return
		}

		// Start a goroutine that will renew the token before it expires.
		go r.tokenRenewer()
	})

	var m sync.Mutex
	r.OnDisconnect(func() {
		m.Lock()
		close(r.disconnect)
		r.disconnect = make(chan bool)
		m.Unlock()
	})

	return r
}

// newRemoteKiteWithClient returns a pointer to new RemoteKite instance.
// The client will be replaced with the given client.
// Used to give the Kite method handler a working RemoteKite to call methods
// on other side.
func (k *Kite) newRemoteKiteWithClient(kite protocol.Kite, auth callAuthentication, client *rpc.Client) *RemoteKite {
	r := k.NewRemoteKite(kite, auth)
	r.client = client
	r.client.Properties()["localKite"] = k
	return r
}

// SetCallTimeout sets the timeout duration for requests made with Call().
func (r *RemoteKite) SetCallTimeout(ms uint) {
	r.callTimeout = time.Duration(ms) * time.Millisecond
}

// Dial connects to the remote Kite. Returns error if it can't.
func (r *RemoteKite) Dial() (err error) {
	addr := r.Kite.Addr()
	r.Log.Info("Dialing remote kite: [%s %s]", r.Kite.Name, addr)
	return r.client.Dial("ws://" + addr + "/dnode")
}

// Dial connects to the remote Kite. If it can't connect, it retries indefinitely.
func (r *RemoteKite) DialForever() {
	addr := r.Kite.Addr()
	r.Log.Info("Dialing remote kite: [%s %s]", r.Kite.Name, addr)
	r.client.DialForever("ws://" + addr + "/dnode")
}

func (r *RemoteKite) Close() {
	r.client.Close()
}

// OnConnect registers a function to run on connect.
func (r *RemoteKite) OnConnect(handler func()) {
	r.client.OnConnect(handler)
}

// OnDisconnect registers a function to run on disconnect.
func (r *RemoteKite) OnDisconnect(handler func()) {
	r.client.OnDisconnect(handler)
}

func (r *RemoteKite) tokenRenewer() {
	for {
		renewTime := r.Authentication.ValidUntil.Add(-30 * time.Second)
		select {
		case <-time.After(renewTime.Sub(time.Now().UTC())):
			if err := r.renewTokenUntilDisconnect(); err != nil {
				return
			}
		case <-r.disconnect:
			return
		}
	}
}

// renewToken retries until the request is successful or disconnect.
func (r *RemoteKite) renewTokenUntilDisconnect() error {
	const retryInterval = 10 * time.Second

	if err := r.renewToken(); err == nil {
		return nil
	}

	for {
		select {
		case <-time.After(retryInterval):
			if err := r.renewToken(); err != nil {
				r.Log.Error("error: %s Cannot renew token for Kite: %s I will retry in %d seconds...", err.Error(), r.Kite.ID, retryInterval)
				continue
			}

			break
		case <-r.disconnect:
			return errors.New("disconnect")
		}
	}

	return nil
}

func (r *RemoteKite) renewToken() error {
	tkn, err := r.localKite.Kontrol.GetToken(&r.Kite)
	if err != nil {
		return err
	}

	validUntil := time.Now().UTC().Add(time.Duration(tkn.TTL) * time.Second)
	r.Authentication.Key = tkn.Key
	r.Authentication.ValidUntil = &validUntil

	return nil
}

// CallOptions is the type of first argument in the dnode message.
// Second argument is a callback function.
// It is used when unmarshalling a dnode message.
type CallOptions struct {
	// Arguments to the method
	WithArgs       *dnode.Partial     `json:"withArgs"`
	Kite           protocol.Kite      `json:"kite"`
	Authentication callAuthentication `json:"authentication"`
}

// callOptionsOut is the same structure with CallOptions.
// It is used when marshalling a dnode message.
type callOptionsOut struct {
	CallOptions
	// Override this when sending because args will not be a *dnode.Partial.
	WithArgs interface{} `json:"withArgs"`
}

// That's what we send as a first argument in dnode message.
func (r *RemoteKite) makeOptions(args interface{}) *callOptionsOut {
	return &callOptionsOut{
		WithArgs: args,
		CallOptions: CallOptions{
			Kite:           r.localKite.Kite,
			Authentication: r.Authentication,
		},
	}
}

type callAuthentication struct {
	// Type can be "kodingKey", "token" or "sessionID" for now.
	Type       string     `json:"type"`
	Key        string     `json:"key"`
	ValidUntil *time.Time `json:"-"`
}

type response struct {
	Result *dnode.Partial
	Err    error
}

// Call makes a blocking method call to the server.
// Waits until the callback function is called by the other side and
// returns the result and the error.
func (r *RemoteKite) Call(method string, args interface{}) (result *dnode.Partial, err error) {
	response := <-r.Go(method, args)
	return response.Result, response.Err
}

// Go makes an unblocking method call to the server.
// It returns a channel that the caller can wait on it to get the response.
func (r *RemoteKite) Go(method string, args interface{}) chan *response {
	// We will return this channel to the caller.
	// It can wait on this channel to get the response.
	r.Log.Debug("Calling method [%s] on kite [%s]", method, r.Name)
	responseChan := make(chan *response, 1)

	r.send(method, args, responseChan)

	return responseChan
}

// send sends the method with callback to the server.
func (r *RemoteKite) send(method string, args interface{}, responseChan chan *response) {
	// To clean the sent callback after response is received.
	// Send/Receive in a channel to prevent race condition because
	// the callback is run in a separate goroutine.
	removeCallback := make(chan uint64, 1)

	// When a callback is called it will send the response to this channel.
	doneChan := make(chan *response, 1)

	opts := r.makeOptions(args)
	cb := r.makeResponseCallback(doneChan, removeCallback)

	callbacks, err := r.client.Call(method, opts, cb)
	if err != nil {
		responseChan <- &response{
			Result: nil,
			Err: fmt.Errorf("Calling method [%s] on [%s] error: %s",
				method, r.Kite.Name, err),
		}
		return
	}

	// Waits until the response has came or the connection has disconnected.
	go func() {
		select {
		case <-r.disconnect:
			responseChan <- &response{nil, errors.New("Client disconnected")}
		case resp := <-doneChan:
			responseChan <- resp
		case <-time.After(r.callTimeout):
			responseChan <- &response{nil, errors.New("Timeout")}

			// Remove the callback function from the map so we do not
			// consume memory for unused callbacks.
			if id, ok := <-removeCallback; ok {
				r.client.RemoveCallback(id)
			}
		}
	}()

	sendCallbackID(callbacks, removeCallback)
}

// sendCallbackID send the callback number to be deleted after response is received.
func sendCallbackID(callbacks map[string]dnode.Path, ch chan uint64) {
	if len(callbacks) > 0 {
		// Find max callback ID.
		max := uint64(0)
		for id, _ := range callbacks {
			i, _ := strconv.ParseUint(id, 10, 64)
			if i > max {
				max = i
			}
		}

		ch <- max
	} else {
		close(ch)
	}
}

// makeResponseCallback prepares and returns a callback function sent to the server.
// The caller of the Call() is blocked until the server calls this callback function.
// Sets theResponse and notifies the caller by sending to done channel.
func (r *RemoteKite) makeResponseCallback(doneChan chan *response, removeCallback <-chan uint64) Callback {
	return Callback(func(request *Request) {
		var (
			err    error          // First argument
			result *dnode.Partial // Second argument
		)

		// Notify that the callback is finished.
		defer func() { doneChan <- &response{result, err} }()

		// Remove the callback function from the map so we do not
		// consume memory for unused callbacks.
		if id, ok := <-removeCallback; ok {
			r.client.RemoveCallback(id)
		}

		// Arguments to our response callback It is a slice of length 2.
		// The first argument is the error string,
		// the second argument is the result.
		responseArgs := request.Args.MustSliceOfLength(2)

		// The second argument is our result.
		result = responseArgs[1]

		// This is error argument. Unmarshal panics if it is null.
		if responseArgs[0] == nil {
			return
		}

		// Read the error argument in response.
		var kiteErr *Error
		err = responseArgs[0].Unmarshal(kiteErr)
		if err != nil {
			return
		}

		err = kiteErr
	})
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"

	r "reflect"

	"strings"

	sci "github.com/samuell/scipipe"
)

func main() {
	p := NewFooToBarReplacer()
	fmt.Println("Process: ", p)
}

type FooToBarReplacer struct {
	sci.Process
	InnerProcess *sci.Process
	InFoo        chan *sci.FileTarget
	OutBar       chan *sci.FileTarget
}

func (p *FooToBarReplacer) Run() {
	// ...
}

func NewFooToBarReplacer() interface{} {
	execFunc := func(task *sci.ShellTask) {
		indata := task.InTargets["foo"].Read()
		indataReplaced := bytes.Replace(indata, []byte("foo"), []byte("bar"), -1)
		task.OutTargets["bar"].WriteTempFile(indataReplaced)
	}
	pathFuncs := map[string]func(*sci.ShellTask) string{
		"bar": func(t *sci.ShellTask) string { return t.InTargets["foo"].GetPath() + ".bar.txt" },
	}
	return NewProcessFromStruct(&FooToBarReplacer{}, execFunc, pathFuncs)
}

func NewProcessFromStruct(procStruct interface{}, execFunc func(*sci.ShellTask), pathFuncs map[string]func(*sci.ShellTask) string) interface{} {
	// Get in-ports of struct
	procStructVal := r.ValueOf(procStruct).Elem()
	for i := 0; i < procStructVal.NumField(); i++ {
		structFieldName := procStructVal.Type().Field(i).Name
		structFieldType := procStructVal.Type().Field(i).Type
		exampleChan := make(chan *sci.FileTarget)
		if strings.HasPrefix(structFieldName, "In") && structFieldType == r.TypeOf(exampleChan) {
			fmt.Println("In-port:", structFieldName)
		} else if strings.HasPrefix(structFieldName, "Out") && structFieldType == r.TypeOf(exampleChan) {
			fmt.Println("Out-port:", structFieldName)
		}
	}
	return procStruct
}
<commit_msg>Collecting in- and out-ports, in example 16<commit_after>package main

import (
	"bytes"
	"fmt"

	r "reflect"

	"strings"

	sci "github.com/samuell/scipipe"
)

func main() {
	p := NewFooToBarReplacer()
	fmt.Println("Process: ", p)
}

// -------------------------------------------
//  Example of defining a new wrapper task
// -------------------------------------------

type FooToBarReplacer struct {
	sci.Process
	InnerProcess *sci.Process
	Run          func(p *FooToBarReplacer)
	InFoo        chan *sci.FileTarget
	OutBar       chan *sci.FileTarget
}

func (p *FooToBarReplacer) Run() {
	// ...
}

func NewFooToBarReplacer() interface{} {
	execFunc := func(task *sci.ShellTask) {
		indata := task.InTargets["foo"].Read()
		indataReplaced := bytes.Replace(indata, []byte("foo"), []byte("bar"), -1)
		task.OutTargets["bar"].WriteTempFile(indataReplaced)
	}
	pathFuncs := map[string]func(*sci.ShellTask) string{
		"bar": func(t *sci.ShellTask) string { return t.InTargets["foo"].GetPath() + ".bar.txt" },
	}
	return NewProcessFromStruct(&FooToBarReplacer{}, execFunc, pathFuncs)
}

// -------------------------------------------
//  New helper methods
// -------------------------------------------

func NewProcessFromStruct(procStruct interface{}, execFunc func(*sci.ShellTask), pathFuncs map[string]func(*sci.ShellTask) string) interface{} {
	// Get in-ports of struct
	inPorts := map[string]chan *sci.FileTarget{}
	outPorts := map[string]chan *sci.FileTarget{}

	procStructVal := r.ValueOf(procStruct).Elem()
	for i := 0; i < procStructVal.NumField(); i++ {
		structFieldName := procStructVal.Type().Field(i).Name
		structFieldType := procStructVal.Type().Field(i).Type
		exampleChan := make(chan *sci.FileTarget)
		if strings.HasPrefix(structFieldName, "In") && structFieldType == r.TypeOf(exampleChan) {
			fmt.Println("In-port:", structFieldName)
			inPorts[strings.ToLower(structFieldName)] = exampleChan // TODO: Change this!
		} else if strings.HasPrefix(structFieldName, "Out") && structFieldType == r.TypeOf(exampleChan) {
			fmt.Println("Out-port:", structFieldName)
			outPorts[strings.ToLower(structFieldName)] = exampleChan // TODO: Change this!
		}
	}
	return procStruct
}
<|endoftext|>


<commit_before>package payment

import (
	"errors"
	"fmt"
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/helpers"
	"sync"
	"time"

	"gopkg.in/fatih/set.v0"
	"gopkg.in/mgo.v2/bson"

	"github.com/stripe/stripe-go"
	"github.com/stripe/stripe-go/customer"
	stripeplan "github.com/stripe/stripe-go/plan"
	"github.com/stripe/stripe-go/sub"
)

var (
	// ErrCustomerNotSubscribedToAnyPlans error for not subscribed users
	ErrCustomerNotSubscribedToAnyPlans = errors.New("user is not subscribed to any plans")
	// ErrCustomerNotExists error for not created users
	ErrCustomerNotExists = errors.New("user is not created for subscription")
	// ErrGroupAlreadyHasSub error when a group tries to create a sub and they try to create another
	ErrGroupAlreadyHasSub = errors.New("group already has a subscription")
)

// Usage holds current usage information, which will be calculated on the fly
type Usage struct {
	User            *UserInfo
	ExpectedPlan    *stripe.Plan
	Due             uint64
	NextBillingDate time.Time
	Subscription    *stripe.Sub
}

// UserInfo holds current info about team's user info
type UserInfo struct {
	Total   int
	Active  int
	Deleted int
}

// DeleteCustomerForGroup deletes the customer for a given group. If customer is
// not registered, returns error. If customer is already deleted, returns success.
func DeleteCustomerForGroup(groupName string) error {
	group, err := modelhelper.GetGroup(groupName)
	if err != nil {
		return err
	}

	if group.Payment.Customer.ID == "" {
		return ErrCustomerNotExists
	}

	if err := deleteCustomer(group.Payment.Customer.ID); err != nil {
		return err
	}

	return modelhelper.UpdateGroupPartial(
		modelhelper.Selector{"_id": group.Id},
		modelhelper.Selector{
			// deleting customer deletes everything belong to that customer in stripe,
			// so say we all
			"$unset": modelhelper.Selector{"payment": ""},
		},
	)
}

// UpdateCustomerForGroup updates customer data of a group`
func UpdateCustomerForGroup(username, groupName string, params *stripe.CustomerParams) (*stripe.Customer, error) {
	group, err := modelhelper.GetGroup(groupName)
	if err != nil {
		return nil, err
	}

	if group.Payment.Customer.ID == "" {
		return nil, ErrCustomerNotExists
	}

	params, err = populateCustomerParams(username, groupName, params)
	if err != nil {
		return nil, err
	}

	return customer.Update(group.Payment.Customer.ID, params)
}

// GetCustomerForGroup get the registered customer info of a group if exists
func GetCustomerForGroup(groupName string) (*stripe.Customer, error) {
	group, err := modelhelper.GetGroup(groupName)
	if err != nil {
		return nil, err
	}

	if group.Payment.Customer.ID == "" {
		return nil, ErrCustomerNotExists
	}

	return customer.Get(group.Payment.Customer.ID, nil)
}

// GetInfoForGroup get the current usage info of a group
func GetInfoForGroup(group *models.Group) (*Usage, error) {
	usage, err := fetchParallelizableeUsageItems(group)
	if err != nil {
		return nil, err
	}

	plan, err := getPlan(usage.Subscription, usage.User.Total)
	if err != nil {
		return nil, err
	}

	usage.ExpectedPlan = plan
	usage.Due = uint64(usage.User.Total) * plan.Amount

	return usage, nil

}

func fetchParallelizableeUsageItems(group *models.Group) (*Usage, error) {
	var infoErr error
	var errMu sync.Mutex
	var wg sync.WaitGroup

	withCheck := func(f func() error) {
		errMu.Lock()
		if infoErr != nil {
			errMu.Unlock()
			return
		}
		errMu.Unlock()

		err := f()
		if err != nil {
			errMu.Lock()
			infoErr = err
			errMu.Unlock()
		}
		wg.Done()
	}

	var cus *stripe.Customer
	wg.Add(1)
	go withCheck(func() (err error) {
		if group.Payment.Customer.ID == "" {
			return ErrCustomerNotExists
		}

		cus, err = customer.Get(group.Payment.Customer.ID, nil)
		return err
	})

	var subscription *stripe.Sub
	wg.Add(1)
	go withCheck(func() (err error) {
		if group.Payment.Subscription.ID == "" {
			return ErrCustomerNotSubscribedToAnyPlans
		}

		subscription, err = sub.Get(group.Payment.Subscription.ID, nil)
		return err
	})

	var activeCount int
	wg.Add(1)
	go withCheck(func() (err error) {
		activeCount, err = calculateActiveUserCount(group.Id)
		return err
	})

	var deletedCount int
	wg.Add(1)
	go withCheck(func() (err error) {
		deletedCount, err = modelhelper.GetDeletedMemberCountByGroupId(group.Id)
		return err
	})

	wg.Wait()

	if infoErr != nil {
		return nil, infoErr
	}

	totalCount := activeCount + deletedCount

	// if the team is in trialing period, and when we want to charge them, do
	// not include the deleted members
	if subscription.Status == "trialing" {
		totalCount = activeCount
	}

	usage := &Usage{
		User: &UserInfo{
			Total:   totalCount,
			Active:  activeCount,
			Deleted: deletedCount,
		},
		ExpectedPlan:    nil,
		Due:             0,
		NextBillingDate: time.Unix(subscription.PeriodEnd, 0),
		Subscription:    subscription,
	}

	return usage, nil
}

func getPlan(subscription *stripe.Sub, totalCount int) (*stripe.Plan, error) {
	plan := subscription.Plan

	if plan.Amount == 0 { // we are on trial period
		return plan, nil
	}

	// in the cases where the active subscription and the have-to-be
	// subscription is different, fetch the real plan from system. This can only
	// happen if the team got more members than the previous subscription's user
	// count in the current month. The subscription will be automatically fixed
	// on the next billing date. We do not change the subscription on each user
	// addition or deletion becasue Stripe charges the user whenever a
	// subscription change happens, so we only change the subscription on the
	// billing date with cancelling the previous subscription & invoice and
	// creating a new subscription with new requirement
	if plan.ID == GetPlanID(totalCount) {
		return plan, nil
	}

	return stripeplan.Get(GetPlanID(totalCount), nil)
}

func createFilter(groupID bson.ObjectId) modelhelper.Selector {
	return modelhelper.Selector{
		"as":         modelhelper.Selector{"$in": []string{"owner", "admin", "member"}},
		"targetName": "JAccount",
		"sourceName": "JGroup",
		"sourceId":   groupID,
	}
}

// calculateActiveUserCount calculates the active user count from the
// relationship collection. I tried using map-reduce first but that was so
// magical from  engineering perspective and we dont have any other usage of it
// in our system. Then i implemented it using "aggregate" framework, that worked
// pretty well indeed but it is fetching all the records from database at once,
// so decided to use our battle tested iter support, it handles iterations, bulk
// operations, timeouts. We are actively using it for iterating over millions of
// records without hardening on the database.
//
//
// Sample aggregate function
//
// db.relationships.aggregate([
//     { "$match": { as: { $in : [ "owner", "admin", "member" ] }, targetName:"JAccount", sourceName:"JGroup", sourceId: ObjectId("") }},
//     // Count all occurrences
//     { "$group": {
//         "_id": {
//             "targetId": "$targetId"
//         },
//         "count": { "$sum": 1 }
//     }},
//     // Sum all occurrences and count distinct
//     { "$group": {
//         "_id": {
//             "targetId": "$_id.targetId"
//         },
//         "totalCount": { "$sum": "$count" },
//         "distinctCount": { "$sum": 1 }
//     }}
// ])
//
func calculateActiveUserCount(groupID bson.ObjectId) (int, error) {
	accounts := set.New()

	iterOptions := helpers.NewIterOptions()
	iterOptions.F = func(rel interface{}) error {
		result, ok := rel.(*models.Relationship)
		if !ok {
			return errors.New("not a relationship")
		}
		accounts.Add(result.TargetId)
		return nil
	}
	iterOptions.CollectionName = modelhelper.RelationshipColl
	iterOptions.Filter = createFilter(groupID)
	iterOptions.Result = &models.Relationship{}
	iterOptions.Limit = 0
	iterOptions.Skip = 0
	// iterOptions.Log = log

	err := helpers.Iter(modelhelper.Mongo, iterOptions)
	if err != nil {
		return 0, err
	}

	return accounts.Size(), nil
}

// CreateCustomerForGroup registers a customer for a group
func CreateCustomerForGroup(username, groupName string, req *stripe.CustomerParams) (*stripe.Customer, error) {
	req, err := populateCustomerParams(username, groupName, req)
	if err != nil {
		return nil, err
	}

	cus, err := customer.New(req)
	if err != nil {
		return nil, err
	}

	group, err := modelhelper.GetGroup(groupName)
	if err != nil {
		return nil, err
	}

	if err := modelhelper.UpdateGroupPartial(
		modelhelper.Selector{"_id": group.Id},
		modelhelper.Selector{
			"$set": modelhelper.Selector{
				"payment.customer.id": cus.ID,
			},
		},
	); err != nil {
		return nil, err
	}

	return cus, nil
}

// deleteCustomer tries to make customer deletion idempotent
func deleteCustomer(customerID string) error {
	cus, err := customer.Del(customerID)
	if cus != nil && cus.Deleted { // if customer is already deleted previously
		return nil
	}

	return err
}

func populateCustomerParams(username, groupName string, initial *stripe.CustomerParams) (*stripe.CustomerParams, error) {
	if initial == nil {
		initial = &stripe.CustomerParams{}
	}

	// whitelisted parameters
	req := &stripe.CustomerParams{
		Token:  initial.Token,
		Coupon: initial.Coupon,
		Source: initial.Source,
		Desc:   initial.Desc,
		Email:  initial.Email,
		Params: initial.Params,
		// plan can not be updated by hand, do not add it to whilelist. It should
		// only be updated automatically on invoice applications
		// Plan: initial.Plan,
	}

	user, err := modelhelper.GetUser(username)
	if err != nil {
		return nil, err
	}

	if req.Desc == "" {
		req.Desc = fmt.Sprintf("%s team", groupName)
	}
	if req.Email == "" {
		req.Email = user.Email
	}

	if req.Params.Meta == nil {
		req.Params.Meta = make(map[string]string)
	}
	req.Params.Meta["groupName"] = groupName
	req.Params.Meta["username"] = username

	return req, nil
}
<commit_msg>payment: fix for missing wg.Done() function call<commit_after>package payment

import (
	"errors"
	"fmt"
	"time"

	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/helpers"

	"gopkg.in/fatih/set.v0"
	"gopkg.in/mgo.v2/bson"

	"github.com/stripe/stripe-go"
	"github.com/stripe/stripe-go/customer"
	stripeplan "github.com/stripe/stripe-go/plan"
	"github.com/stripe/stripe-go/sub"
	"golang.org/x/sync/errgroup"
)

var (
	// ErrCustomerNotSubscribedToAnyPlans error for not subscribed users
	ErrCustomerNotSubscribedToAnyPlans = errors.New("user is not subscribed to any plans")
	// ErrCustomerNotExists error for not created users
	ErrCustomerNotExists = errors.New("user is not created for subscription")
	// ErrGroupAlreadyHasSub error when a group tries to create a sub and they try to create another
	ErrGroupAlreadyHasSub = errors.New("group already has a subscription")
)

// Usage holds current usage information, which will be calculated on the fly
type Usage struct {
	User            *UserInfo
	ExpectedPlan    *stripe.Plan
	Due             uint64
	NextBillingDate time.Time
	Subscription    *stripe.Sub
}

// UserInfo holds current info about team's user info
type UserInfo struct {
	Total   int
	Active  int
	Deleted int
}

// DeleteCustomerForGroup deletes the customer for a given group. If customer is
// not registered, returns error. If customer is already deleted, returns success.
func DeleteCustomerForGroup(groupName string) error {
	group, err := modelhelper.GetGroup(groupName)
	if err != nil {
		return err
	}

	if group.Payment.Customer.ID == "" {
		return ErrCustomerNotExists
	}

	if err := deleteCustomer(group.Payment.Customer.ID); err != nil {
		return err
	}

	return modelhelper.UpdateGroupPartial(
		modelhelper.Selector{"_id": group.Id},
		modelhelper.Selector{
			// deleting customer deletes everything belong to that customer in stripe,
			// so say we all
			"$unset": modelhelper.Selector{"payment": ""},
		},
	)
}

// UpdateCustomerForGroup updates customer data of a group`
func UpdateCustomerForGroup(username, groupName string, params *stripe.CustomerParams) (*stripe.Customer, error) {
	group, err := modelhelper.GetGroup(groupName)
	if err != nil {
		return nil, err
	}

	if group.Payment.Customer.ID == "" {
		return nil, ErrCustomerNotExists
	}

	params, err = populateCustomerParams(username, groupName, params)
	if err != nil {
		return nil, err
	}

	return customer.Update(group.Payment.Customer.ID, params)
}

// GetCustomerForGroup get the registered customer info of a group if exists
func GetCustomerForGroup(groupName string) (*stripe.Customer, error) {
	group, err := modelhelper.GetGroup(groupName)
	if err != nil {
		return nil, err
	}

	if group.Payment.Customer.ID == "" {
		return nil, ErrCustomerNotExists
	}

	return customer.Get(group.Payment.Customer.ID, nil)
}

// GetInfoForGroup get the current usage info of a group
func GetInfoForGroup(group *models.Group) (*Usage, error) {
	usage, err := fetchParallelizableeUsageItems(group)
	if err != nil {
		return nil, err
	}

	plan, err := getPlan(usage.Subscription, usage.User.Total)
	if err != nil {
		return nil, err
	}

	usage.ExpectedPlan = plan
	usage.Due = uint64(usage.User.Total) * plan.Amount

	return usage, nil

}

func fetchParallelizableeUsageItems(group *models.Group) (*Usage, error) {
	var g errgroup.Group

	// Stripe customer.
	var cus *stripe.Customer
	g.Go(func() (err error) {
		if group.Payment.Customer.ID == "" {
			return ErrCustomerNotExists
		}

		cus, err = customer.Get(group.Payment.Customer.ID, nil)
		return err
	})

	// Stripe subscription.
	var subscription *stripe.Sub
	g.Go(func() (err error) {
		if group.Payment.Subscription.ID == "" {
			return ErrCustomerNotSubscribedToAnyPlans
		}

		subscription, err = sub.Get(group.Payment.Subscription.ID, nil)
		return err
	})

	// Active users count.
	var activeCount int
	g.Go(func() (err error) {
		activeCount, err = calculateActiveUserCount(group.Id)
		return err
	})

	// Deleted members count.
	var deletedCount int
	g.Go(func() (err error) {
		deletedCount, err = modelhelper.GetDeletedMemberCountByGroupId(group.Id)
		return err
	})

	if err := g.Wait(); err != nil {
		return nil, err
	}

	totalCount := activeCount + deletedCount

	// if the team is in trialing period, and when we want to charge them, do
	// not include the deleted members
	if subscription.Status == "trialing" {
		totalCount = activeCount
	}

	usage := &Usage{
		User: &UserInfo{
			Total:   totalCount,
			Active:  activeCount,
			Deleted: deletedCount,
		},
		ExpectedPlan:    nil,
		Due:             0,
		NextBillingDate: time.Unix(subscription.PeriodEnd, 0),
		Subscription:    subscription,
	}

	return usage, nil
}

func getPlan(subscription *stripe.Sub, totalCount int) (*stripe.Plan, error) {
	plan := subscription.Plan

	if plan.Amount == 0 { // we are on trial period
		return plan, nil
	}

	// in the cases where the active subscription and the have-to-be
	// subscription is different, fetch the real plan from system. This can only
	// happen if the team got more members than the previous subscription's user
	// count in the current month. The subscription will be automatically fixed
	// on the next billing date. We do not change the subscription on each user
	// addition or deletion becasue Stripe charges the user whenever a
	// subscription change happens, so we only change the subscription on the
	// billing date with cancelling the previous subscription & invoice and
	// creating a new subscription with new requirement
	if plan.ID == GetPlanID(totalCount) {
		return plan, nil
	}

	return stripeplan.Get(GetPlanID(totalCount), nil)
}

func createFilter(groupID bson.ObjectId) modelhelper.Selector {
	return modelhelper.Selector{
		"as":         modelhelper.Selector{"$in": []string{"owner", "admin", "member"}},
		"targetName": "JAccount",
		"sourceName": "JGroup",
		"sourceId":   groupID,
	}
}

// calculateActiveUserCount calculates the active user count from the
// relationship collection. I tried using map-reduce first but that was so
// magical from  engineering perspective and we dont have any other usage of it
// in our system. Then i implemented it using "aggregate" framework, that worked
// pretty well indeed but it is fetching all the records from database at once,
// so decided to use our battle tested iter support, it handles iterations, bulk
// operations, timeouts. We are actively using it for iterating over millions of
// records without hardening on the database.
//
//
// Sample aggregate function
//
// db.relationships.aggregate([
//     { "$match": { as: { $in : [ "owner", "admin", "member" ] }, targetName:"JAccount", sourceName:"JGroup", sourceId: ObjectId("") }},
//     // Count all occurrences
//     { "$group": {
//         "_id": {
//             "targetId": "$targetId"
//         },
//         "count": { "$sum": 1 }
//     }},
//     // Sum all occurrences and count distinct
//     { "$group": {
//         "_id": {
//             "targetId": "$_id.targetId"
//         },
//         "totalCount": { "$sum": "$count" },
//         "distinctCount": { "$sum": 1 }
//     }}
// ])
//
func calculateActiveUserCount(groupID bson.ObjectId) (int, error) {
	accounts := set.New()

	iterOptions := helpers.NewIterOptions()
	iterOptions.F = func(rel interface{}) error {
		result, ok := rel.(*models.Relationship)
		if !ok {
			return errors.New("not a relationship")
		}
		accounts.Add(result.TargetId)
		return nil
	}
	iterOptions.CollectionName = modelhelper.RelationshipColl
	iterOptions.Filter = createFilter(groupID)
	iterOptions.Result = &models.Relationship{}
	iterOptions.Limit = 0
	iterOptions.Skip = 0
	// iterOptions.Log = log

	err := helpers.Iter(modelhelper.Mongo, iterOptions)
	if err != nil {
		return 0, err
	}

	return accounts.Size(), nil
}

// CreateCustomerForGroup registers a customer for a group
func CreateCustomerForGroup(username, groupName string, req *stripe.CustomerParams) (*stripe.Customer, error) {
	req, err := populateCustomerParams(username, groupName, req)
	if err != nil {
		return nil, err
	}

	cus, err := customer.New(req)
	if err != nil {
		return nil, err
	}

	group, err := modelhelper.GetGroup(groupName)
	if err != nil {
		return nil, err
	}

	if err := modelhelper.UpdateGroupPartial(
		modelhelper.Selector{"_id": group.Id},
		modelhelper.Selector{
			"$set": modelhelper.Selector{
				"payment.customer.id": cus.ID,
			},
		},
	); err != nil {
		return nil, err
	}

	return cus, nil
}

// deleteCustomer tries to make customer deletion idempotent
func deleteCustomer(customerID string) error {
	cus, err := customer.Del(customerID)
	if cus != nil && cus.Deleted { // if customer is already deleted previously
		return nil
	}

	return err
}

func populateCustomerParams(username, groupName string, initial *stripe.CustomerParams) (*stripe.CustomerParams, error) {
	if initial == nil {
		initial = &stripe.CustomerParams{}
	}

	// whitelisted parameters
	req := &stripe.CustomerParams{
		Token:  initial.Token,
		Coupon: initial.Coupon,
		Source: initial.Source,
		Desc:   initial.Desc,
		Email:  initial.Email,
		Params: initial.Params,
		// plan can not be updated by hand, do not add it to whilelist. It should
		// only be updated automatically on invoice applications
		// Plan: initial.Plan,
	}

	user, err := modelhelper.GetUser(username)
	if err != nil {
		return nil, err
	}

	if req.Desc == "" {
		req.Desc = fmt.Sprintf("%s team", groupName)
	}
	if req.Email == "" {
		req.Email = user.Email
	}

	if req.Params.Meta == nil {
		req.Params.Meta = make(map[string]string)
	}
	req.Params.Meta["groupName"] = groupName
	req.Params.Meta["username"] = username

	return req, nil
}
<|endoftext|>


<commit_before>/* Go (cgo) interface to libgeoip */
package geoip

/*
#cgo CFLAGS: -I/opt/local/include -I/usr/local/include -I/usr/include
#cgo LDFLAGS: -lGeoIP -L/opt/local/lib -L/usr/local/lib -L/usr/lib
#include <stdio.h>
#include <errno.h>
#include <GeoIP.h>

//typedef GeoIP* GeoIP_pnt
*/
import "C"

import (
	"log"
	"os"
	"unsafe"
)

// Set Quiet to true to silence warnings when the database can't be opened
// (though libgeoip still prints to errout)
var Quiet bool = false

type GeoIP struct {
	db *C.GeoIP
}

// Open opens a GeoIP database, all formats supported by libgeoip are supported though
// there are only functions to access the country information in this API.
// The database is opened in MEMORY_CACHE mode, if you need to optimize for memory
// instead of performance you should change this.
// If you don't pass a filename, it will try opening the database from
// a list of common paths.
func Open(files ...string) *GeoIP {
	if len(files) == 0 {
		files = []string{
			"/usr/share/GeoIP/GeoIP.dat",       // Linux default
			"/usr/share/local/GeoIP/GeoIP.dat", // source install?
			"/usr/local/share/GeoIP/GeoIP.dat", // FreeBSD
			"/opt/local/share/GeoIP/GeoIP.dat", // MacPorts
			"/usr/share/GeoIP/GeoIP.dat",       // ArchLinux
		}
	}

	g := &GeoIP{}

	var err error

	for _, file := range files {

		// libgeoip prints errors if it can't open the file, so check first
		if _, err := os.Stat(file); err != nil {
			if os.IsExist(err) {
				log.Println(err)
			}
			continue
		}

		cbase := C.CString(file)
		g.db, err = C.GeoIP_open(cbase, C.GEOIP_MEMORY_CACHE)
		C.free(unsafe.Pointer(cbase))
		if g.db != nil && err != nil {
			break
		}
	}
	if err != nil && !Quiet {
		log.Println("Error opening GeoIP database", files, err)
	}

	if g.db != nil {
		C.GeoIP_set_charset(g.db, C.GEOIP_CHARSET_UTF8)
		return g
	}

	return nil
}

// GetOrg takes an IPv4 address string and returns the org name for that IP.
// Requires the geoip org database
func (gi *GeoIP) GetOrg(ip string) string {
	if gi.db == nil {
		return ""
	}
	cip := C.CString(ip)
	cname := C.GeoIP_name_by_addr(gi.db, cip)
	C.free(unsafe.Pointer(cip))
	if cname != nil {
		rets := C.GoString(cname)
		return rets
	}
	return ""
}

// GetCountry takes an IPv4 address string and returns the country code for that IP.
func (gi *GeoIP) GetCountry(ip string) string {
	if gi.db == nil {
		return ""
	}
	cip := C.CString(ip)
	ccountry := C.GeoIP_country_code_by_addr(gi.db, cip)
	C.free(unsafe.Pointer(cip))
	if ccountry != nil {
		rets := C.GoString(ccountry)
		return rets
	}
	return ""
}

// GetCountry_v6 works the same as GetCountry except for IPv6 addresses, be sure to
// load a database with IPv6 data to get any results.
func (gi *GeoIP) GetCountry_v6(ip string) string {
	if gi.db == nil {
		return ""
	}
	cip := C.CString(ip)
	ccountry := C.GeoIP_country_code_by_addr_v6(gi.db, cip)
	C.free(unsafe.Pointer(cip))
	if ccountry != nil {
		rets := C.GoString(ccountry)
		return rets
	}
	return ""
}
<commit_msg>Close GeoIP database when it goes out of scope<commit_after>/* Go (cgo) interface to libgeoip */
package geoip

/*
#cgo CFLAGS: -I/opt/local/include -I/usr/local/include -I/usr/include
#cgo LDFLAGS: -lGeoIP -L/opt/local/lib -L/usr/local/lib -L/usr/lib
#include <stdio.h>
#include <errno.h>
#include <GeoIP.h>

//typedef GeoIP* GeoIP_pnt
*/
import "C"

import (
	"log"
	"os"
	"runtime"
	"unsafe"
)

// Set Quiet to true to silence warnings when the database can't be opened
// (though libgeoip still prints to errout)
var Quiet bool = false

type GeoIP struct {
	db *C.GeoIP
}

func (gi *GeoIP) Free() {
	if gi == nil {
		return
	}
	if gi.db == nil {
		gi = nil
		return
	}
	C.GeoIP_delete(gi.db)
	gi = nil
	return
}

// Open opens a GeoIP database, all formats supported by libgeoip are supported though
// there are only functions to access the country information in this API.
// The database is opened in MEMORY_CACHE mode, if you need to optimize for memory
// instead of performance you should change this.
// If you don't pass a filename, it will try opening the database from
// a list of common paths.
func Open(files ...string) *GeoIP {
	if len(files) == 0 {
		files = []string{
			"/usr/share/GeoIP/GeoIP.dat",       // Linux default
			"/usr/share/local/GeoIP/GeoIP.dat", // source install?
			"/usr/local/share/GeoIP/GeoIP.dat", // FreeBSD
			"/opt/local/share/GeoIP/GeoIP.dat", // MacPorts
			"/usr/share/GeoIP/GeoIP.dat",       // ArchLinux
		}
	}

	g := &GeoIP{}
	runtime.SetFinalizer(g, (*GeoIP).Free)

	var err error

	for _, file := range files {

		// libgeoip prints errors if it can't open the file, so check first
		if _, err := os.Stat(file); err != nil {
			if os.IsExist(err) {
				log.Println(err)
			}
			continue
		}

		cbase := C.CString(file)
		g.db, err = C.GeoIP_open(cbase, C.GEOIP_MEMORY_CACHE)
		C.free(unsafe.Pointer(cbase))
		if g.db != nil && err != nil {
			break
		}
	}
	if err != nil && !Quiet {
		log.Println("Error opening GeoIP database", files, err)
	}

	if g.db != nil {
		C.GeoIP_set_charset(g.db, C.GEOIP_CHARSET_UTF8)
		return g
	}

	return nil
}

// GetOrg takes an IPv4 address string and returns the org name for that IP.
// Requires the geoip org database
func (gi *GeoIP) GetOrg(ip string) string {
	if gi.db == nil {
		return ""
	}
	cip := C.CString(ip)
	cname := C.GeoIP_name_by_addr(gi.db, cip)
	C.free(unsafe.Pointer(cip))
	if cname != nil {
		rets := C.GoString(cname)
		return rets
	}
	return ""
}

// GetCountry takes an IPv4 address string and returns the country code for that IP.
func (gi *GeoIP) GetCountry(ip string) string {
	if gi.db == nil {
		return ""
	}
	cip := C.CString(ip)
	ccountry := C.GeoIP_country_code_by_addr(gi.db, cip)
	C.free(unsafe.Pointer(cip))
	if ccountry != nil {
		rets := C.GoString(ccountry)
		return rets
	}
	return ""
}

// GetCountry_v6 works the same as GetCountry except for IPv6 addresses, be sure to
// load a database with IPv6 data to get any results.
func (gi *GeoIP) GetCountry_v6(ip string) string {
	if gi.db == nil {
		return ""
	}
	cip := C.CString(ip)
	ccountry := C.GeoIP_country_code_by_addr_v6(gi.db, cip)
	C.free(unsafe.Pointer(cip))
	if ccountry != nil {
		rets := C.GoString(ccountry)
		return rets
	}
	return ""
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/davecgh/go-spew/spew"
	"os"
	"syscall"
)

var daemonName = "npserver"

func main() {

	initFlags()

	// check permissions
	if os.Getuid() != 0 && os.Geteuid() != 0 {
		fmt.Printf("npserver-daemon should be run as root, have uid=%d and euid=%d\n", os.Getuid(), os.Geteuid())
		os.Exit(1)
	}

	// check start || sto
	if flags.Start == flags.Stop {
		fmt.Println("need --start or --stop flag")
	}

	if flags.Start {
		startDaemon()
	}
	if flags.Stop {
		stopDaemon()
	}

	// all good :)
}

func startDaemon() {
	// setup args for daemon call
	args := []string{
		fmt.Sprintf("--name=%s", daemonName),
		"--noconfig",
		"--errlog=/var/log/npserver-daemon.log",
		"--output=/var/log/npserver.log",
		fmt.Sprintf("--pidfile=%s", flags.PIDFile),
		"--unsafe",
		"--",
		"/usr/local/bin/npserver",
	}

	// append extra args to args
	args = append(args, extraArgs...)

	// start process
	proc, err := os.StartProcess("/usr/bin/daemon", args, &os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Sys: &syscall.SysProcAttr{
			Credential: &syscall.Credential{
				Uid: uint32(os.Geteuid()),
				Gid: uint32(os.Getegid()),
			},
		},
	})
	if err != nil {
		fmt.Printf("os/exec returned an error: '%s'\n", err)
		os.Exit(1)
	}

	// wait for daemon to be ready
	_, err = proc.Wait()
	if err != nil {
		fmt.Printf("proc.Wait() failed. %s\n", err)
		os.Exit(1)
	}
}

func stopDaemon() {
	// stopping arguments
	args := []string{
		"--stop",
		fmt.Sprintf("--name=%s", daemonName),
		fmt.Sprintf("--pidfile=%s", flags.PIDFile),
	}
	spew.Dump(args)

	// start process
	proc, err := os.StartProcess("/usr/bin/daemon", args, &os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Sys: &syscall.SysProcAttr{
			Credential: &syscall.Credential{
				Uid: uint32(os.Geteuid()),
				Gid: uint32(os.Getegid()),
			},
		},
	})
	if err != nil {
		fmt.Printf("os/exec returned an error: '%s'\n", err)
		os.Exit(1)
	}

	// wait for daemon to be ready
	_, err = proc.Wait()
	if err != nil {
		fmt.Printf("proc.Wait() failed. %s\n", err)
		os.Exit(1)
	}
	// // open pid file (if available)
	// pidFile, err := os.Open(flags.PIDFile)
	// if err != nil {
	// 	if os.IsNotExist(err) {
	// 		fmt.Printf("it looks like npserver is not running")
	// 		os.Exit(0)
	// 	}
	// 	fmt.Printf("error on opening pidfile: %s\n", err)
	// 	os.Exit(1)
	// }

	// // read all file contents
	// pidFileContents, err := ioutil.ReadAll(pidFile)
	// pidFile.Close()
	// if err != nil {
	// 	fmt.Printf("error reading pidfile contents: %s\n", err)
	// 	os.Exit(1)
	// }

	// pidFileContentsString := string(pidFileContents)

	// // strip eventual whitespace
	// pidFileContentsString = strings.TrimRight(pidFileContentsString, " \r\n\t")

	// // convert pid string to pid int
	// pid, err := strconv.Atoi(pidFileContentsString)
	// if err != nil {
	// 	fmt.Printf("error parsing pidfile contents: %s\n", err)
	// 	os.Exit(1)
	// }

	// // lookup process
	// proc, err := os.FindProcess(pid)
	// if err != nil {
	// 	fmt.Printf("error finding process with pid %d: %s\n", pid, err)
	// 	os.Exit(1)
	// }

	// // signal process to stop
	// err = proc.Signal(os.Interrupt)
	// if err != nil {
	// 	fmt.Printf("error sending interrupt signal to npserver: %s\n", err)
	// 	os.Exit(1)
	// }

	// // wait until process is done
	// state, err := proc.Wait()
	// if err != nil {
	// 	if err.Error() != "wait: no child processes" {
	// 		fmt.Printf("error waiting for process to stop: %s\n", err)
	// 		os.Exit(1)
	// 	}
	// } else {
	// 	if !state.Exited() || !state.Success() {
	// 		fmt.Printf("npserver process exited badly")
	// 		os.Exit(1)
	// 	}
	// }

	// // remove pid file
	// err = os.Remove(flags.PIDFile)
	// if err != nil {
	// 	fmt.Printf("error removing pid file: %s\n", err)
	// }
}
<commit_msg>killall...<commit_after>package main

import (
	"fmt"
	"github.com/davecgh/go-spew/spew"
	"os"
	"os/exec"
	"syscall"
)

var daemonName = "npserver"

func main() {

	initFlags()

	// check permissions
	if os.Getuid() != 0 && os.Geteuid() != 0 {
		fmt.Printf("npserver-daemon should be run as root, have uid=%d and euid=%d\n", os.Getuid(), os.Geteuid())
		os.Exit(1)
	}

	// check start || sto
	if flags.Start == flags.Stop {
		fmt.Println("need --start or --stop flag")
	}

	if flags.Start {
		startDaemon()
	}
	if flags.Stop {
		stopDaemon()
	}

	// all good :)
}

func startDaemon() {
	// setup args for daemon call
	args := []string{
		fmt.Sprintf("--name=%s", daemonName),
		"--noconfig",
		"--errlog=/var/log/npserver-daemon.log",
		"--output=/var/log/npserver.log",
		fmt.Sprintf("--pidfile=%s", flags.PIDFile),
		"--unsafe",
		"--",
		"/usr/local/bin/npserver",
	}

	// append extra args to args
	args = append(args, extraArgs...)

	// start process
	proc, err := os.StartProcess("/usr/bin/daemon", args, &os.ProcAttr{
		Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
		Sys: &syscall.SysProcAttr{
			Credential: &syscall.Credential{
				Uid: uint32(os.Geteuid()),
				Gid: uint32(os.Getegid()),
			},
		},
	})
	if err != nil {
		fmt.Printf("os/exec returned an error: '%s'\n", err)
		os.Exit(1)
	}

	// wait for daemon to be ready
	_, err = proc.Wait()
	if err != nil {
		fmt.Printf("proc.Wait() failed. %s\n", err)
		os.Exit(1)
	}
}

func stopDaemon() {
	// possibly racy.
	// no guarantee that npserver cleaned up before new npserver is being copied/started
	cmd := exec.Command("killall", "-s", "SIGINT", "npserver")
	cmd.Run()

	// // stopping arguments
	// args := []string{
	// 	"--stop",
	// 	fmt.Sprintf("--name=%s", daemonName),
	// 	fmt.Sprintf("--pidfile=%s", flags.PIDFile),
	// }
	// spew.Dump(args)

	// // start process
	// proc, err := os.StartProcess("/usr/bin/daemon", args, &os.ProcAttr{
	// 	Files: []*os.File{os.Stdin, os.Stdout, os.Stderr},
	// 	Sys: &syscall.SysProcAttr{
	// 		Credential: &syscall.Credential{
	// 			Uid: uint32(os.Geteuid()),
	// 			Gid: uint32(os.Getegid()),
	// 		},
	// 	},
	// })
	// if err != nil {
	// 	fmt.Printf("os/exec returned an error: '%s'\n", err)
	// 	os.Exit(1)
	// }

	// // wait for daemon to be ready
	// _, err = proc.Wait()
	// if err != nil {
	// 	fmt.Printf("proc.Wait() failed. %s\n", err)
	// 	os.Exit(1)
	// }

	// // open pid file (if available)
	// pidFile, err := os.Open(flags.PIDFile)
	// if err != nil {
	// 	if os.IsNotExist(err) {
	// 		fmt.Printf("it looks like npserver is not running")
	// 		os.Exit(0)
	// 	}
	// 	fmt.Printf("error on opening pidfile: %s\n", err)
	// 	os.Exit(1)
	// }

	// // read all file contents
	// pidFileContents, err := ioutil.ReadAll(pidFile)
	// pidFile.Close()
	// if err != nil {
	// 	fmt.Printf("error reading pidfile contents: %s\n", err)
	// 	os.Exit(1)
	// }

	// pidFileContentsString := string(pidFileContents)

	// // strip eventual whitespace
	// pidFileContentsString = strings.TrimRight(pidFileContentsString, " \r\n\t")

	// // convert pid string to pid int
	// pid, err := strconv.Atoi(pidFileContentsString)
	// if err != nil {
	// 	fmt.Printf("error parsing pidfile contents: %s\n", err)
	// 	os.Exit(1)
	// }

	// // lookup process
	// proc, err := os.FindProcess(pid)
	// if err != nil {
	// 	fmt.Printf("error finding process with pid %d: %s\n", pid, err)
	// 	os.Exit(1)
	// }

	// // signal process to stop
	// err = proc.Signal(os.Interrupt)
	// if err != nil {
	// 	fmt.Printf("error sending interrupt signal to npserver: %s\n", err)
	// 	os.Exit(1)
	// }

	// // wait until process is done
	// state, err := proc.Wait()
	// if err != nil {
	// 	if err.Error() != "wait: no child processes" {
	// 		fmt.Printf("error waiting for process to stop: %s\n", err)
	// 		os.Exit(1)
	// 	}
	// } else {
	// 	if !state.Exited() || !state.Success() {
	// 		fmt.Printf("npserver process exited badly")
	// 		os.Exit(1)
	// 	}
	// }

	// // remove pid file
	// err = os.Remove(flags.PIDFile)
	// if err != nil {
	// 	fmt.Printf("error removing pid file: %s\n", err)
	// }
}
<|endoftext|>


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

// Package flags provides an extensive command line option parser.
// The flags package is similar in functionality to the go builtin flag package
// but provides more options and uses reflection to provide a convenient and
// succinct way of specifying command line options.
//
// Supported features:
//     Options with short names (-v)
//     Options with long names (--verbose)
//     Options with and without arguments (bool v.s. other type)
//     Options with optional arguments and default values
//     Multiple option groups each containing a set of options
//     Generate and print well-formatted help message
//     Passing remaining command line arguments after -- (optional)
//     Ignoring unknown command line options (optional)
//     Supports -I/usr/include -I=/usr/include -I /usr/include option argument specification
//     Supports multiple short options -aux
//     Supports all primitive go types (string, int{8..64}, uint{8..64}, float)
//     Supports same option multiple times (can store in slice or last option counts)
//     Supports maps
//     Supports function callbacks
//
// Additional features specific to Windows:
//     Options with short names (/v)
//     Options with long names (/verbose)
//     Windows-style options with arguments use a colon as the delimiter
//     Modify generated help message with Windows-style / options
//
// The flags package uses structs, reflection and struct field tags
// to allow users to specify command line options. This results in very simple
// and consise specification of your application options. For example:
//
//     type Options struct {
//         Verbose []bool `short:"v" long:"verbose" description:"Show verbose debug information"`
//     }
//
// This specifies one option with a short name -v and a long name --verbose.
// When either -v or --verbose is found on the command line, a 'true' value
// will be appended to the Verbose field. e.g. when specifying -vvv, the
// resulting value of Verbose will be {[true, true, true]}.
//
// Slice options work exactly the same as primitive type options, except that
// whenever the option is encountered, a value is appended to the slice.
//
// Map options from string to primitive type are also supported. On the command
// line, you specify the value for such an option as key:value. For example
//
//     type Options struct {
//         AuthorInfo string[string] `short:"a"`
//     }
//
// Then, the AuthorInfo map can be filled with something like
// -a name:Jesse -a "surname:van den Kieboom".
//
// Finally, for full control over the conversion between command line argument
// values and options, user defined types can choose to implement the Marshaler
// and Unmarshaler interfaces.
//
// Available field tags:
//     short:          the short name of the option (single character)
//     long:           the long name of the option
//     description:    the description of the option (optional)
//     optional:       whether an argument of the option is optional (optional)
//     optional-value: the value of an optional option when the option occurs
//                     without an argument. This tag can be specified multiple
//                     times in the case of maps or slices (optional)
//     default:        the default value of an option. This tag can be specified
//                     multiple times in the case of slices or maps (optional).
//     default-mask:   when specified, this value will be displayed in the help
//                     instead of the actual default value. This is useful
//                     mostly for hiding otherwise sensitive information from
//                     showing up in the help. If default-mask takes the special
//                     value "-", then no default value will be shown at all
//                     (optional)
//     required:       whether an option is required to appear on the command
//                     line. If a required option is not present, the parser
//                     will return ErrRequired.
//     base:           a base (radix) used to convert strings to integer values,
//                     the default base is 10 (i.e. decimal) (optional)
//     value-name:     the name of the argument value (to be shown in the help,
//                     (optional)
//     group:          when specified on a struct field, makes the struct field
//                     a separate group with the given name (optional).
//     command:        when specified on a struct field, makes the struct field
//                     a (sub)command with the given name (optional).
//
//     subcommands-optional: when specified on a command struct field, makes
//                           any subcommands of that command optional.
//
// Either short: or long: must be specified to make the field eligible as an
// option.
//
//
// Option groups:
//
// Option groups are a simple way to semantically separate your options. The
// only real difference is in how your options will appear in the builtin
// generated help. All options in a particular group are shown together in the
// help under the name of the group.
//
// There are currently three ways to specify option groups.
//
//     1. Use NewNamedParser specifying the various option groups.
//     2. Use AddGroup to add a group to an existing parser.
//     3. Add a struct field to the toplevel options annotated with the
//        group:"group-name" tag.
//
//
//
// Commands:
//
// The flags package also has basic support for commands. Commands are often
// used in monolithic applications that support various commands or actions.
// Take git for example, all of the add, commit, checkout, etc. are called
// commands. Using commands you can easily separate multiple functions of your
// application.
//
// There are currently two ways to specifiy a command.
//
//     1. Use AddCommand on an existing parser.
//     2. Add a struct field to your options struct annotated with the
//        command:"command-name" tag.
//
// The most common, idiomatic way to implement commands is to define a global
// parser instance and implement each command in a separate file. These
// command files should define a go init function which calls AddCommand on
// the global parser.
//
// When parsing ends and there is an active command and that command implements
// the Commander interface, then its Execute method will be run with the
// remaining command line arguments.
//
// Command structs can have options which become valid to parse after the
// command has been specified on the command line. It is currently not valid
// to specify options from the parent level of the command after the command
// name has occurred. Thus, given a toplevel option "-v" and a command "add":
//
//     Valid:   ./app -v add
//     Invalid: ./app add -v
//
package flags
<commit_msg>group field tags<commit_after>// Copyright 2012 Jesse van den Kieboom. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package flags provides an extensive command line option parser.
// The flags package is similar in functionality to the go builtin flag package
// but provides more options and uses reflection to provide a convenient and
// succinct way of specifying command line options.
//
// Supported features:
//     Options with short names (-v)
//     Options with long names (--verbose)
//     Options with and without arguments (bool v.s. other type)
//     Options with optional arguments and default values
//     Multiple option groups each containing a set of options
//     Generate and print well-formatted help message
//     Passing remaining command line arguments after -- (optional)
//     Ignoring unknown command line options (optional)
//     Supports -I/usr/include -I=/usr/include -I /usr/include option argument specification
//     Supports multiple short options -aux
//     Supports all primitive go types (string, int{8..64}, uint{8..64}, float)
//     Supports same option multiple times (can store in slice or last option counts)
//     Supports maps
//     Supports function callbacks
//
// Additional features specific to Windows:
//     Options with short names (/v)
//     Options with long names (/verbose)
//     Windows-style options with arguments use a colon as the delimiter
//     Modify generated help message with Windows-style / options
//
// The flags package uses structs, reflection and struct field tags
// to allow users to specify command line options. This results in very simple
// and consise specification of your application options. For example:
//
//     type Options struct {
//         Verbose []bool `short:"v" long:"verbose" description:"Show verbose debug information"`
//     }
//
// This specifies one option with a short name -v and a long name --verbose.
// When either -v or --verbose is found on the command line, a 'true' value
// will be appended to the Verbose field. e.g. when specifying -vvv, the
// resulting value of Verbose will be {[true, true, true]}.
//
// Slice options work exactly the same as primitive type options, except that
// whenever the option is encountered, a value is appended to the slice.
//
// Map options from string to primitive type are also supported. On the command
// line, you specify the value for such an option as key:value. For example
//
//     type Options struct {
//         AuthorInfo string[string] `short:"a"`
//     }
//
// Then, the AuthorInfo map can be filled with something like
// -a name:Jesse -a "surname:van den Kieboom".
//
// Finally, for full control over the conversion between command line argument
// values and options, user defined types can choose to implement the Marshaler
// and Unmarshaler interfaces.
//
// Available field tags:
//     short:       the short name of the option (single character)
//     long:        the long name of the option
//     required:    whether an option is required to appear on the command line.
//                  If a required option is not present, the parser will return
//                  ErrRequired.
//     description: the description of the option (optional)
//
//     optional:       whether an argument of the option is optional (optional)
//     optional-value: the value of an optional option when the option occurs
//                     without an argument. This tag can be specified multiple
//                     times in the case of maps or slices (optional)
//     default:        the default value of an option. This tag can be specified
//                     multiple times in the case of slices or maps (optional).
//     default-mask:   when specified, this value will be displayed in the help
//                     instead of the actual default value. This is useful
//                     mostly for hiding otherwise sensitive information from
//                     showing up in the help. If default-mask takes the special
//                     value "-", then no default value will be shown at all
//                     (optional)
//     value-name:     the name of the argument value (to be shown in the help,
//                     (optional)
//
//     base: a base (radix) used to convert strings to integer values, the default
//           base is 10 (i.e. decimal) (optional)
//
//     group:                when specified on a struct field, makes the struct
//                           field a separate group with the given name (optional).
//     command:              when specified on a struct field, makes the struct
//                           field a (sub)command with the given name (optional).
//     subcommands-optional: when specified on a command struct field, makes
//                           any subcommands of that command optional.
//
// Either short: or long: must be specified to make the field eligible as an
// option.
//
//
// Option groups:
//
// Option groups are a simple way to semantically separate your options. The
// only real difference is in how your options will appear in the builtin
// generated help. All options in a particular group are shown together in the
// help under the name of the group.
//
// There are currently three ways to specify option groups.
//
//     1. Use NewNamedParser specifying the various option groups.
//     2. Use AddGroup to add a group to an existing parser.
//     3. Add a struct field to the toplevel options annotated with the
//        group:"group-name" tag.
//
//
//
// Commands:
//
// The flags package also has basic support for commands. Commands are often
// used in monolithic applications that support various commands or actions.
// Take git for example, all of the add, commit, checkout, etc. are called
// commands. Using commands you can easily separate multiple functions of your
// application.
//
// There are currently two ways to specifiy a command.
//
//     1. Use AddCommand on an existing parser.
//     2. Add a struct field to your options struct annotated with the
//        command:"command-name" tag.
//
// The most common, idiomatic way to implement commands is to define a global
// parser instance and implement each command in a separate file. These
// command files should define a go init function which calls AddCommand on
// the global parser.
//
// When parsing ends and there is an active command and that command implements
// the Commander interface, then its Execute method will be run with the
// remaining command line arguments.
//
// Command structs can have options which become valid to parse after the
// command has been specified on the command line. It is currently not valid
// to specify options from the parent level of the command after the command
// name has occurred. Thus, given a toplevel option "-v" and a command "add":
//
//     Valid:   ./app -v add
//     Invalid: ./app add -v
//
package flags
<|endoftext|>


<commit_before>package runtime

import (
	"testing"
	"time"
	"github.com/devfeel/dotweb/test"
	"sync"
)

const (
	DefaultTestGCInterval = 2

	TEST_CACHE_KEY = "joe"
	TEST_CACHE_VALUE = "zou"
	//int value
	TEST_CACHE_INT_VALUE = 1

	//int64 value
	TEST_CACHE_INT64_VALUE = 1
)

func TestRuntimeCache_Get(t *testing.T) {
	cache:=NewRuntimeCache()
	cache.Set(TEST_CACHE_KEY,TEST_CACHE_VALUE,5)
	//check value
	go func(cache *RuntimeCache,t *testing.T) {
		time.Sleep(4*time.Second)
		value,err:=cache.Get(TEST_CACHE_KEY)

		test.Nil(t,err)
		test.Equal(t,TEST_CACHE_VALUE,value)
	}(cache,t)

	//check expired
	go func(cache *RuntimeCache,t *testing.T) {
		time.Sleep(5*time.Second)
		value,err:=cache.Exists(TEST_CACHE_KEY)

		test.Nil(t,err)
		test.Equal(t,true,value)
	}(cache,t)

	time.Sleep(5*time.Second)
}

func TestRuntimeCache_GetInt(t *testing.T) {
	testRuntimeCache(t,TEST_CACHE_INT_VALUE,func(cache *RuntimeCache,key string)(interface{}, error){
		return cache.GetInt(key)
	})
}

func TestRuntimeCache_GetInt64(t *testing.T) {
	testRuntimeCache(t,TEST_CACHE_INT64_VALUE,func(cache *RuntimeCache,key string)(interface{}, error){
		return cache.GetInt64(key)
	})
}

func TestRuntimeCache_GetString(t *testing.T) {
	testRuntimeCache(t,TEST_CACHE_VALUE,func(cache *RuntimeCache,key string)(interface{}, error){
		return cache.GetString(key)
	})
}

func testRuntimeCache(t *testing.T,insertValue interface{},f func(cache *RuntimeCache,key string)(interface{}, error)) {
	cache:=NewRuntimeCache()
	cache.Set(TEST_CACHE_KEY,insertValue,5)
	//check value
	go func(cache *RuntimeCache,t *testing.T) {
		time.Sleep(4*time.Second)
		value,err:=f(cache,TEST_CACHE_KEY)

		test.Nil(t,err)
		test.Equal(t,insertValue,value)
	}(cache,t)

	time.Sleep(5*time.Second)
}

func TestRuntimeCache_Delete(t *testing.T) {
	cache:=NewRuntimeCache()
	cache.Set(TEST_CACHE_KEY,TEST_CACHE_VALUE,5)

	value,e:=cache.Get(TEST_CACHE_KEY)

	test.Nil(t,e)
	test.Equal(t,TEST_CACHE_VALUE,value)

	cache.Delete(TEST_CACHE_KEY)

	value,e=cache.Get(TEST_CACHE_KEY)
	test.Nil(t,e)
	test.Nil(t,value)
}

func TestRuntimeCache_ClearAll(t *testing.T) {
	cache:=NewRuntimeCache()
	cache.Set(TEST_CACHE_KEY,TEST_CACHE_VALUE,5)
	cache.Set("2",TEST_CACHE_VALUE,5)
	cache.Set("3",TEST_CACHE_VALUE,5)

	val_2, err := cache.GetString("2")
	if err != nil{
		t.Error(err)
	}
	test.Equal(t,TEST_CACHE_VALUE, val_2)

	cache.ClearAll()
	exists_2, err := cache.Exists("2")
	if err != nil{
		t.Error(err)
	}
	if exists_2{
		t.Error("exists 2 but need not exists")
	}
}

func TestRuntimeCache_Incr(t *testing.T) {
	cache:=NewRuntimeCache()
	var wg sync.WaitGroup
	wg.Add(2)

	go func(cache *RuntimeCache) {
		for i := 0; i < 50; i++ {
			cache.Incr(TEST_CACHE_KEY)
		}

		wg.Add(-1)
	}(cache)

	go func(cache *RuntimeCache) {
		for i := 0; i < 50; i++ {
			cache.Incr(TEST_CACHE_KEY)
		}
		wg.Add(-1)
	}(cache)

	wg.Wait()

	value,e:=cache.GetInt(TEST_CACHE_KEY)
	test.Nil(t,e)

	test.Equal(t,100,value)
}

func TestRuntimeCache_Decr(t *testing.T) {
	cache:=NewRuntimeCache()
	var wg sync.WaitGroup
	wg.Add(2)

	go func(cache *RuntimeCache) {
		for i := 0; i < 50; i++ {
			cache.Decr(TEST_CACHE_KEY)
		}

		wg.Add(-1)
	}(cache)

	go func(cache *RuntimeCache) {
		for i := 0; i < 50; i++ {
			cache.Decr(TEST_CACHE_KEY)
		}
		wg.Add(-1)
	}(cache)

	wg.Wait()

	value,e:=cache.GetInt(TEST_CACHE_KEY)
	test.Nil(t,e)

	test.Equal(t,-100,value)
}<commit_msg>#### Version 1.4.9.4 * 调整: dotweb.Classic移除自动UseRequestLog逻辑 * 调整：Session Redis模式时，新增sessionReExpire用于重置有效期，避免调用SessionUpdate导致不必要的redis数据交换 * 调整：Cache.Runtime调整map为sync.Map * 状态页：访问/dotweb/state时增加CurrentTime字段输出 * BUG： 修正Reponse自动释放时未重置body字段，导致内存溢出BUG * 2018-05-10 13:30<commit_after>package runtime

import (
	"testing"
	"time"
	"github.com/devfeel/dotweb/test"
	"sync"
)

const (
	DefaultTestGCInterval = 2

	TEST_CACHE_KEY = "joe"
	TEST_CACHE_VALUE = "zou"
	//int value
	TEST_CACHE_INT_VALUE = 1

	//int64 value
	TEST_CACHE_INT64_VALUE = 1
)

func TestRuntimeCache_Get(t *testing.T) {
	cache:=NewRuntimeCache()
	cache.Set(TEST_CACHE_KEY,TEST_CACHE_VALUE,5)
	//check value
	go func(cache *RuntimeCache,t *testing.T) {
		time.Sleep(4*time.Second)
		value,err:=cache.Get(TEST_CACHE_KEY)

		test.Nil(t,err)
		test.Equal(t,TEST_CACHE_VALUE,value)
	}(cache,t)

	//check expired
	go func(cache *RuntimeCache,t *testing.T) {
		time.Sleep(5*time.Second)
		value,err:=cache.Exists(TEST_CACHE_KEY)

		test.Nil(t,err)
		test.Equal(t,true,value)
	}(cache,t)

	time.Sleep(5*time.Second)
}

func TestRuntimeCache_GetInt(t *testing.T) {
	testRuntimeCache(t,TEST_CACHE_INT_VALUE,func(cache *RuntimeCache,key string)(interface{}, error){
		return cache.GetInt(key)
	})
}

func TestRuntimeCache_GetInt64(t *testing.T) {
	testRuntimeCache(t,TEST_CACHE_INT64_VALUE,func(cache *RuntimeCache,key string)(interface{}, error){
		return cache.GetInt64(key)
	})
}

func TestRuntimeCache_GetString(t *testing.T) {
	testRuntimeCache(t,TEST_CACHE_VALUE,func(cache *RuntimeCache,key string)(interface{}, error){
		return cache.GetString(key)
	})
}

func testRuntimeCache(t *testing.T,insertValue interface{},f func(cache *RuntimeCache,key string)(interface{}, error)) {
	cache:=NewRuntimeCache()
	cache.Set(TEST_CACHE_KEY,insertValue,5)
	//check value
	go func(cache *RuntimeCache,t *testing.T) {
		time.Sleep(4*time.Second)
		value,err:=f(cache,TEST_CACHE_KEY)

		test.Nil(t,err)
		test.Equal(t,insertValue,value)
	}(cache,t)

	time.Sleep(5*time.Second)
}

func TestRuntimeCache_Delete(t *testing.T) {
	cache:=NewRuntimeCache()
	cache.Set(TEST_CACHE_KEY,TEST_CACHE_VALUE,5)

	value,e:=cache.Get(TEST_CACHE_KEY)

	test.Nil(t,e)
	test.Equal(t,TEST_CACHE_VALUE,value)

	cache.Delete(TEST_CACHE_KEY)

	value,e=cache.Get(TEST_CACHE_KEY)
	test.Nil(t,e)
	test.Nil(t,value)
}

func TestRuntimeCache_ClearAll(t *testing.T) {
	cache:=NewRuntimeCache()
	cache.Set(TEST_CACHE_KEY,TEST_CACHE_VALUE,5)
	cache.Set("2",TEST_CACHE_VALUE,5)
	cache.Set("3",TEST_CACHE_VALUE,5)

	val2, err := cache.GetString("2")
	if err != nil{
		t.Error(err)
	}
	test.Equal(t,TEST_CACHE_VALUE, val2)

	cache.ClearAll()
	exists2, err := cache.Exists("2")
	if err != nil{
		t.Error(err)
	}
	if exists2{
		t.Error("exists 2 but need not exists")
	}
}

func TestRuntimeCache_Incr(t *testing.T) {
	cache:=NewRuntimeCache()
	var wg sync.WaitGroup
	wg.Add(2)

	go func(cache *RuntimeCache) {
		for i := 0; i < 50; i++ {
			cache.Incr(TEST_CACHE_KEY)
		}

		wg.Add(-1)
	}(cache)

	go func(cache *RuntimeCache) {
		for i := 0; i < 50; i++ {
			cache.Incr(TEST_CACHE_KEY)
		}
		wg.Add(-1)
	}(cache)

	wg.Wait()

	value,e:=cache.GetInt(TEST_CACHE_KEY)
	test.Nil(t,e)

	test.Equal(t,100,value)
}

func TestRuntimeCache_Decr(t *testing.T) {
	cache:=NewRuntimeCache()
	var wg sync.WaitGroup
	wg.Add(2)

	go func(cache *RuntimeCache) {
		for i := 0; i < 50; i++ {
			cache.Decr(TEST_CACHE_KEY)
		}

		wg.Add(-1)
	}(cache)

	go func(cache *RuntimeCache) {
		for i := 0; i < 50; i++ {
			cache.Decr(TEST_CACHE_KEY)
		}
		wg.Add(-1)
	}(cache)

	wg.Wait()

	value,e:=cache.GetInt(TEST_CACHE_KEY)
	test.Nil(t,e)

	test.Equal(t,-100,value)
}<|endoftext|>