<commit_before>package apply

import (
	"context"
	"fmt"

	"k8s.io/client-go/tools/cache"

	"kubevirt.io/kubevirt/pkg/controller"

	rbacv1 "k8s.io/api/rbac/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/client-go/log"
	"kubevirt.io/kubevirt/pkg/virt-operator/resource/generate/rbac"
)

type RoleType int

const (
	TypeRole               RoleType = iota
	TypeClusterRole        RoleType = iota
	TypeRoleBinding        RoleType = iota
	TypeClusterRoleBinding RoleType = iota
)

func (r *Reconciler) createOrUpdateClusterRole(cr *rbacv1.ClusterRole, imageTag string, imageRegistry string, id string) error {
	return r.createOrUpdate(cr, imageTag, imageRegistry, id, TypeClusterRole, false)
}

func (r *Reconciler) createOrUpdateClusterRoleBinding(crb *rbacv1.ClusterRoleBinding, imageTag string, imageRegistry string, id string) error {
	return r.createOrUpdate(crb, imageTag, imageRegistry, id, TypeClusterRoleBinding, false)
}

func (r *Reconciler) createOrUpdateRole(role *rbacv1.Role, imageTag string, imageRegistry string, id string) error {
	return r.createOrUpdate(role, imageTag, imageRegistry, id, TypeRole, true)
}

func (r *Reconciler) createOrUpdateRoleBinding(rb *rbacv1.RoleBinding, imageTag string, imageRegistry string, id string) error {
	return r.createOrUpdate(rb, imageTag, imageRegistry, id, TypeRoleBinding, true)
}

func (r *Reconciler) createOrUpdate(role interface{},
	imageTag, imageRegistry, id string,
	roleType RoleType,
	avoidIfServiceAccount bool) (err error) {

	roleTypeName := getRoleTypeName(roleType)
	createRole := r.getRoleCreateFunction(role, roleType)
	updateRole := r.getRoleUpdateFunction(role, roleType)

	cachedRole, exists, _ := r.getRoleCache(roleType).Get(role)
	roleMeta := getRoleMetaObject(role, roleType)
	if avoidIfServiceAccount && !r.stores.ServiceMonitorEnabled && (roleMeta.Name == rbac.MONITOR_SERVICEACCOUNT_NAME) {
		return nil
	}

	injectOperatorMetadata(r.kv, roleMeta, imageTag, imageRegistry, id, true)
	if !exists {
		// Create non existent
		err = createRole()
		if err != nil {
			return fmt.Errorf("unable to create %v %+v: %v", roleTypeName, role, err)
		}
		log.Log.V(2).Infof("%v %v created", roleTypeName, roleMeta.GetName())
	} else if !objectMatchesVersion(getRoleMetaObject(cachedRole, roleType), imageTag, imageRegistry, id, r.kv.GetGeneration()) {
		// Update existing, we don't need to patch for rbac rules.
		err = updateRole()
		if err != nil {
			return fmt.Errorf("unable to update %v %+v: %v", roleTypeName, role, err)
		}
		log.Log.V(2).Infof("%v %v updated", roleTypeName, roleMeta.GetName())

	} else {
		log.Log.V(4).Infof("%v %v already exists", roleTypeName, roleMeta.GetName())
	}

	return nil
}

func (r *Reconciler) getRoleCreateFunction(obj interface{}, roleType RoleType) (createFunc func() error) {

	rbacObj := r.clientset.RbacV1()
	namespace := r.kv.Namespace

	raiseExpectation := func(exp *controller.UIDTrackingControllerExpectations) {
		exp.RaiseExpectations(r.kvKey, 1, 0)
	}
	lowerExpectationIfErr := func(exp *controller.UIDTrackingControllerExpectations, err error) {
		if err != nil {
			exp.LowerExpectations(r.kvKey, 1, 0)
		}
	}

	switch roleType {
	case TypeRole:
		role := obj.(*rbacv1.Role)

		createFunc = func() error {
			raiseExpectation(r.expectations.Role)
			_, err := rbacObj.Roles(namespace).Create(context.Background(), role, metav1.CreateOptions{})
			lowerExpectationIfErr(r.expectations.Role, err)
			return err
		}
	case TypeClusterRole:
		role := obj.(*rbacv1.ClusterRole)

		createFunc = func() error {
			raiseExpectation(r.expectations.ClusterRole)
			_, err := rbacObj.ClusterRoles().Create(context.Background(), role, metav1.CreateOptions{})
			lowerExpectationIfErr(r.expectations.ClusterRole, err)
			return err
		}
	case TypeRoleBinding:
		roleBinding := obj.(*rbacv1.RoleBinding)

		createFunc = func() error {
			raiseExpectation(r.expectations.RoleBinding)
			_, err := rbacObj.RoleBindings(namespace).Create(context.Background(), roleBinding, metav1.CreateOptions{})
			lowerExpectationIfErr(r.expectations.RoleBinding, err)
			return err
		}
	case TypeClusterRoleBinding:
		roleBinding := obj.(*rbacv1.ClusterRoleBinding)

		createFunc = func() error {
			raiseExpectation(r.expectations.ClusterRoleBinding)
			_, err := rbacObj.ClusterRoleBindings().Create(context.Background(), roleBinding, metav1.CreateOptions{})
			lowerExpectationIfErr(r.expectations.ClusterRoleBinding, err)
			return err
		}
	}

	return
}

func (r *Reconciler) getRoleUpdateFunction(obj interface{}, roleType RoleType) (updateFunc func() (err error)) {
	rbacObj := r.clientset.RbacV1()
	namespace := r.kv.Namespace

	switch roleType {
	case TypeRole:
		role := obj.(*rbacv1.Role)

		updateFunc = func() (err error) {
			_, err = rbacObj.Roles(namespace).Update(context.Background(), role, metav1.UpdateOptions{})
			return err
		}
	case TypeClusterRole:
		role := obj.(*rbacv1.ClusterRole)

		updateFunc = func() (err error) {
			_, err = rbacObj.ClusterRoles().Update(context.Background(), role, metav1.UpdateOptions{})
			return err
		}
	case TypeRoleBinding:
		roleBinding := obj.(*rbacv1.RoleBinding)

		updateFunc = func() (err error) {
			_, err = rbacObj.RoleBindings(namespace).Update(context.Background(), roleBinding, metav1.UpdateOptions{})
			return err
		}
	case TypeClusterRoleBinding:
		roleBinding := obj.(*rbacv1.ClusterRoleBinding)

		updateFunc = func() (err error) {
			_, err = rbacObj.ClusterRoleBindings().Update(context.Background(), roleBinding, metav1.UpdateOptions{})
			return err
		}
	}

	return
}

func getRoleTypeName(roleType RoleType) (name string) {
	switch roleType {
	case TypeRole:
		name = "role"
	case TypeClusterRole:
		name = "clusterrole"
	case TypeRoleBinding:
		name = "rolebinding"
	case TypeClusterRoleBinding:
		name = "clusterrolebinding"
	}

	return
}

func getRoleMetaObject(role interface{}, roleType RoleType) (meta *metav1.ObjectMeta) {
	switch roleType {
	case TypeRole:
		role := role.(*rbacv1.Role)
		meta = &role.ObjectMeta
	case TypeClusterRole:
		role := role.(*rbacv1.ClusterRole)
		meta = &role.ObjectMeta
	case TypeRoleBinding:
		roleBinding := role.(*rbacv1.RoleBinding)
		meta = &roleBinding.ObjectMeta
	case TypeClusterRoleBinding:
		roleBinding := role.(*rbacv1.ClusterRoleBinding)
		meta = &roleBinding.ObjectMeta
	}

	return
}

func (r *Reconciler) getRoleCache(roleType RoleType) (cache cache.Store) {
	switch roleType {
	case TypeRole:
		cache = r.stores.RoleCache
	case TypeClusterRole:
		cache = r.stores.ClusterRoleCache
	case TypeRoleBinding:
		cache = r.stores.RoleBindingCache
	case TypeClusterRoleBinding:
		cache = r.stores.ClusterRoleBindingCache
	}

	return cache
}
<commit_msg>Reconcile Role and ClusterRole resources<commit_after>package apply

import (
	"context"
	"fmt"

	"github.com/openshift/library-go/pkg/operator/resource/resourcemerge"
	"k8s.io/client-go/tools/cache"

	"kubevirt.io/kubevirt/pkg/controller"

	rbacv1 "k8s.io/api/rbac/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"kubevirt.io/client-go/log"
	"kubevirt.io/kubevirt/pkg/virt-operator/resource/generate/rbac"
)

type RoleType int

const (
	TypeRole               RoleType = iota
	TypeClusterRole        RoleType = iota
	TypeRoleBinding        RoleType = iota
	TypeClusterRoleBinding RoleType = iota
)

func (r *Reconciler) createOrUpdateClusterRole(cr *rbacv1.ClusterRole, imageTag string, imageRegistry string, id string) error {
	return r.createOrUpdate(cr, imageTag, imageRegistry, id, TypeClusterRole, false)
}

func (r *Reconciler) createOrUpdateClusterRoleBinding(crb *rbacv1.ClusterRoleBinding, imageTag string, imageRegistry string, id string) error {
	return r.createOrUpdate(crb, imageTag, imageRegistry, id, TypeClusterRoleBinding, false)
}

func (r *Reconciler) createOrUpdateRole(role *rbacv1.Role, imageTag string, imageRegistry string, id string) error {
	return r.createOrUpdate(role, imageTag, imageRegistry, id, TypeRole, true)
}

func (r *Reconciler) createOrUpdateRoleBinding(rb *rbacv1.RoleBinding, imageTag string, imageRegistry string, id string) error {
	return r.createOrUpdate(rb, imageTag, imageRegistry, id, TypeRoleBinding, true)
}

func (r *Reconciler) createOrUpdate(role interface{},
	imageTag, imageRegistry, id string,
	roleType RoleType,
	avoidIfServiceAccount bool) (err error) {

	roleTypeName := getRoleTypeName(roleType)
	createRole := r.getRoleCreateFunction(role, roleType)
	updateRole := r.getRoleUpdateFunction(role, roleType)

	cachedRole, exists, _ := r.getRoleCache(roleType).Get(role)
	roleMeta := getRoleMetaObject(role, roleType)
	if avoidIfServiceAccount && !r.stores.ServiceMonitorEnabled && (roleMeta.Name == rbac.MONITOR_SERVICEACCOUNT_NAME) {
		return nil
	}

	injectOperatorMetadata(r.kv, roleMeta, imageTag, imageRegistry, id, true)
	if !exists {
		// Create non existent
		err = createRole()
		if err != nil {
			return fmt.Errorf("unable to create %v %+v: %v", roleTypeName, role, err)
		}
		log.Log.V(2).Infof("%v %v created", roleTypeName, roleMeta.GetName())
		return nil
	}

	modified := resourcemerge.BoolPtr(false)
	cachedRoleMeta := getRoleMetaObject(cachedRole, roleType)
	resourcemerge.EnsureObjectMeta(modified, cachedRoleMeta.DeepCopy(), *roleMeta)

	// there was no change to metadata, the generation matched
	if !*modified && areRoleRulesEqual(role, cachedRole, roleType) {
		log.Log.V(4).Infof("%v %v already exists", roleTypeName, roleMeta.GetName())
		return nil
	}

	// Update existing, we don't need to patch for rbac rules.
	err = updateRole()
	if err != nil {
		return fmt.Errorf("unable to update %v %+v: %v", roleTypeName, role, err)
	}
	log.Log.V(2).Infof("%v %v updated", roleTypeName, roleMeta.GetName())

	return nil
}

func (r *Reconciler) getRoleCreateFunction(obj interface{}, roleType RoleType) (createFunc func() error) {

	rbacObj := r.clientset.RbacV1()
	namespace := r.kv.Namespace

	raiseExpectation := func(exp *controller.UIDTrackingControllerExpectations) {
		exp.RaiseExpectations(r.kvKey, 1, 0)
	}
	lowerExpectationIfErr := func(exp *controller.UIDTrackingControllerExpectations, err error) {
		if err != nil {
			exp.LowerExpectations(r.kvKey, 1, 0)
		}
	}

	switch roleType {
	case TypeRole:
		role := obj.(*rbacv1.Role)

		createFunc = func() error {
			raiseExpectation(r.expectations.Role)
			_, err := rbacObj.Roles(namespace).Create(context.Background(), role, metav1.CreateOptions{})
			lowerExpectationIfErr(r.expectations.Role, err)
			return err
		}
	case TypeClusterRole:
		role := obj.(*rbacv1.ClusterRole)

		createFunc = func() error {
			raiseExpectation(r.expectations.ClusterRole)
			_, err := rbacObj.ClusterRoles().Create(context.Background(), role, metav1.CreateOptions{})
			lowerExpectationIfErr(r.expectations.ClusterRole, err)
			return err
		}
	case TypeRoleBinding:
		roleBinding := obj.(*rbacv1.RoleBinding)

		createFunc = func() error {
			raiseExpectation(r.expectations.RoleBinding)
			_, err := rbacObj.RoleBindings(namespace).Create(context.Background(), roleBinding, metav1.CreateOptions{})
			lowerExpectationIfErr(r.expectations.RoleBinding, err)
			return err
		}
	case TypeClusterRoleBinding:
		roleBinding := obj.(*rbacv1.ClusterRoleBinding)

		createFunc = func() error {
			raiseExpectation(r.expectations.ClusterRoleBinding)
			_, err := rbacObj.ClusterRoleBindings().Create(context.Background(), roleBinding, metav1.CreateOptions{})
			lowerExpectationIfErr(r.expectations.ClusterRoleBinding, err)
			return err
		}
	}

	return
}

func (r *Reconciler) getRoleUpdateFunction(obj interface{}, roleType RoleType) (updateFunc func() (err error)) {
	rbacObj := r.clientset.RbacV1()
	namespace := r.kv.Namespace

	switch roleType {
	case TypeRole:
		role := obj.(*rbacv1.Role)

		updateFunc = func() (err error) {
			_, err = rbacObj.Roles(namespace).Update(context.Background(), role, metav1.UpdateOptions{})
			return err
		}
	case TypeClusterRole:
		role := obj.(*rbacv1.ClusterRole)

		updateFunc = func() (err error) {
			_, err = rbacObj.ClusterRoles().Update(context.Background(), role, metav1.UpdateOptions{})
			return err
		}
	case TypeRoleBinding:
		roleBinding := obj.(*rbacv1.RoleBinding)

		updateFunc = func() (err error) {
			_, err = rbacObj.RoleBindings(namespace).Update(context.Background(), roleBinding, metav1.UpdateOptions{})
			return err
		}
	case TypeClusterRoleBinding:
		roleBinding := obj.(*rbacv1.ClusterRoleBinding)

		updateFunc = func() (err error) {
			_, err = rbacObj.ClusterRoleBindings().Update(context.Background(), roleBinding, metav1.UpdateOptions{})
			return err
		}
	}

	return
}

func getRoleTypeName(roleType RoleType) (name string) {
	switch roleType {
	case TypeRole:
		name = "role"
	case TypeClusterRole:
		name = "clusterrole"
	case TypeRoleBinding:
		name = "rolebinding"
	case TypeClusterRoleBinding:
		name = "clusterrolebinding"
	}

	return
}

func getRoleMetaObject(role interface{}, roleType RoleType) (meta *metav1.ObjectMeta) {
	switch roleType {
	case TypeRole:
		role := role.(*rbacv1.Role)
		meta = &role.ObjectMeta
	case TypeClusterRole:
		role := role.(*rbacv1.ClusterRole)
		meta = &role.ObjectMeta
	case TypeRoleBinding:
		roleBinding := role.(*rbacv1.RoleBinding)
		meta = &roleBinding.ObjectMeta
	case TypeClusterRoleBinding:
		roleBinding := role.(*rbacv1.ClusterRoleBinding)
		meta = &roleBinding.ObjectMeta
	}

	return
}

func areRoleRulesEqual(role1 interface{}, role2 interface{}, roleType RoleType) (equal bool) {
	// This is to avoid using reflections for performance reasons
	arePolicyRulesEqual := func(pr1 []rbacv1.PolicyRule, pr2 []rbacv1.PolicyRule) bool {
		if len(pr1) != len(pr2) {
			return false
		}

		areStringListsEqual := func(strList1 []string, strList2 []string) bool {
			if len(strList1) != len(strList2) {
				return false
			}
			for i := range strList1 {
				if strList1[i] != strList2[i] {
					return false
				}
			}
			return true
		}

		for i := range pr1 {
			if !areStringListsEqual(pr1[i].Verbs, pr2[i].Verbs) || !areStringListsEqual(pr1[i].Resources, pr2[i].Resources) ||
				!areStringListsEqual(pr1[i].APIGroups, pr2[i].APIGroups) || !areStringListsEqual(pr1[i].NonResourceURLs, pr2[i].NonResourceURLs) ||
				!areStringListsEqual(pr1[i].ResourceNames, pr2[i].ResourceNames) {
				return false
			}
		}

		return true
	}

	switch roleType {
	case TypeRole:
		role1Obj := role1.(*rbacv1.Role)
		role2Obj := role2.(*rbacv1.Role)
		equal = arePolicyRulesEqual(role1Obj.Rules, role2Obj.Rules)
	case TypeClusterRole:
		role1Obj := role1.(*rbacv1.ClusterRole)
		role2Obj := role2.(*rbacv1.ClusterRole)
		equal = arePolicyRulesEqual(role1Obj.Rules, role2Obj.Rules)
	// Bindings do not have "rules" attribute
	case TypeRoleBinding:
		fallthrough
	case TypeClusterRoleBinding:
		equal = true
	}

	return
}

func (r *Reconciler) getRoleCache(roleType RoleType) (cache cache.Store) {
	switch roleType {
	case TypeRole:
		cache = r.stores.RoleCache
	case TypeClusterRole:
		cache = r.stores.ClusterRoleCache
	case TypeRoleBinding:
		cache = r.stores.RoleBindingCache
	case TypeClusterRoleBinding:
		cache = r.stores.ClusterRoleBindingCache
	}

	return cache
}
<|endoftext|>


<commit_before>package responses

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/micro/cli"
	"github.com/perthgophers/puddle/messagerouter"
	"github.com/pkg/errors"
)

func PortfolioWrapper(cr *messagerouter.CommandRequest, w messagerouter.ResponseWriter) error {
	cmd := cli.NewApp()
	cmd.Name = "Portfolio Management"
	cmd.Usage = "Add, track and manage your crypto portfolio!"
	cmd.Action = func(c *cli.Context) {
		fmt.Println(c.Args())
	}
	cmd.Commands = []cli.Command{
		{
			Name:   "register, r",
			Usage:  "Register yourself for portfolio tracking",
			Action: PortfolioRegister,
		},
		{
			Name:  "get, g",
			Usage: "Get your current portoflio",
			Action: func(c *cli.Context) {
				err := Portfolio(cr, w)
				if err != nil {
					w.WriteError(err.Error())
					return
				}
			},
		},
	}
	msg := strings.Fields(cr.Text)
	err := cmd.Run(msg)
	if err != nil {
		return err
	}
	return nil
}

func PortfolioRegister(c *cli.Context) {
	fmt.Println("Called")
}

// Portfolio will return a string containing the current rate in USD
func Portfolio(cr *messagerouter.CommandRequest, w messagerouter.ResponseWriter) error {

	if cr.Username != "nii236" {
		return errors.New("invalid username " + cr.Username)
	}
	ethAmt := os.Getenv("ETH")
	btcAmt := os.Getenv("BTC")
	bchAmt := os.Getenv("BCH")

	if ethAmt == "" || btcAmt == "" || bchAmt == "" {
		return errors.New("portfolio not provided in environment variables")
	}
	type HTTPResponse struct {
		Last string `json:"last"`
	}

	resp, err := http.Get("https://api.coinmarketcap.com/v1/ticker/")
	if err != nil {
		return err
	}
	result := &Tickers{}
	err = json.NewDecoder(resp.Body).Decode(result)
	if err != nil {
		return err
	}

	ETHticker := result.Symbol("ETH")
	BTCticker := result.Symbol("BTC")
	BCHticker := result.Symbol("BCH")

	if ETHticker == nil || BTCticker == nil || BCHticker == nil {
		return errors.New("could not find tickers from API response")
	}

	ethTotal, err := calcValue(ethAmt, ETHticker.PriceUsd)
	if err != nil {
		return errors.Wrap(err, "could not calculate total value")
	}
	btcTotal, err := calcValue(btcAmt, BTCticker.PriceUsd)
	if err != nil {
		return errors.Wrap(err, "could not calculate total value")
	}
	bchTotal, err := calcValue(bchAmt, BCHticker.PriceUsd)
	if err != nil {
		return errors.Wrap(err, "could not calculate total value")
	}

	total := ethTotal + btcTotal + bchTotal
	aud, err := usdToAud(total)
	if err != nil {
		return errors.Wrap(err, "could not convert to AUD")
	}
	w.Write(fmt.Sprintf("Your crypto net worth is: %.2f USD (%.2f AUD)", total, aud))
	return nil
}

func calcValue(ownAmt, tickerAmt string) (float64, error) {
	own, err := strconv.ParseFloat(ownAmt, 64)
	if err != nil {
		return 0, err
	}
	tick, err := strconv.ParseFloat(tickerAmt, 64)
	if err != nil {
		return 0, err
	}
	return own * tick, err
}

type forexTicker struct {
	Base  string `json:"base"`
	Date  string `json:"date"`
	Rates struct {
		AUD float64 `json:"AUD"`
	} `json:"rates"`
}

func usdToAud(usd float64) (float64, error) {

	resp, err := http.Get("https://api.fixer.io/latest?symbols=USD,AUD")
	if err != nil {
		return 0, err
	}
	result := &forexTicker{}
	err = json.NewDecoder(resp.Body).Decode(result)
	if err != nil {
		return 0, err
	}

	converted := result.Rates.AUD * usd
	return converted, nil
}

func init() {
	Handle("!portfolio", PortfolioWrapper)
}
<commit_msg>basic cli adaptor<commit_after>package responses

import (
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/micro/cli"
	"github.com/perthgophers/puddle/messagerouter"
	"github.com/pkg/errors"
)

type PortfolioCommand struct {
	*cli.App
}

func NewPortfolioCommand() *PortfolioCommand {
	return &PortfolioCommand{}
}

func (app *PortfolioCommand) Initialise(cr *messagerouter.CommandRequest, w messagerouter.ResponseWriter) {

	cmd := cli.NewApp()
	cmd.Name = "Portfolio Management"
	cmd.Usage = "Add, track and manage your crypto portfolio!"
	cmd.Commands = []cli.Command{
		{
			Name:    "register",
			Aliases: []string{"r"},
			Usage:   "Register yourself for portfolio tracking",
			Action:  PortfolioRegister(cr, w),
		},
		{
			Name:    "get",
			Aliases: []string{"g"},
			Usage:   "Get your current portoflio",
			Action:  PortfolioGet(cr, w),
		},
	}

	app.App = cmd
}

// PortfolioRegister registers the user for portfolio tracking
func PortfolioRegister(cr *messagerouter.CommandRequest, w messagerouter.ResponseWriter) func(c *cli.Context) {
	return func(c *cli.Context) {
		fmt.Println("Register")
		return
	}
}

// PortfolioGet gets the current user's portfolio
func PortfolioGet(cr *messagerouter.CommandRequest, w messagerouter.ResponseWriter) func(c *cli.Context) {
	return func(c *cli.Context) {
		if cr.Username != "nii236" {
			w.WriteError(errors.New("invalid username " + cr.Username).Error())
			return
		}

		ethAmt := os.Getenv("ETH")
		btcAmt := os.Getenv("BTC")
		bchAmt := os.Getenv("BCH")

		if ethAmt == "" || btcAmt == "" || bchAmt == "" {
			w.WriteError(errors.New("portfolio not provided in environment variables").Error())
			return
		}
		type HTTPResponse struct {
			Last string `json:"last"`
		}

		resp, err := http.Get("https://api.coinmarketcap.com/v1/ticker/")
		if err != nil {
			w.WriteError(err.Error())
		}
		result := &Tickers{}
		err = json.NewDecoder(resp.Body).Decode(result)
		if err != nil {
			w.WriteError(err.Error())
		}

		ETHticker := result.Symbol("ETH")
		BTCticker := result.Symbol("BTC")
		BCHticker := result.Symbol("BCH")

		if ETHticker == nil || BTCticker == nil || BCHticker == nil {
			w.WriteError(errors.New("could not find tickers from API response").Error())
			return
		}

		ethTotal, err := calcValue(ethAmt, ETHticker.PriceUsd)
		if err != nil {
			w.WriteError(errors.Wrap(err, "could not calculate total value").Error())

		}
		btcTotal, err := calcValue(btcAmt, BTCticker.PriceUsd)
		if err != nil {
			w.WriteError(errors.Wrap(err, "could not calculate total value").Error())

		}
		bchTotal, err := calcValue(bchAmt, BCHticker.PriceUsd)
		if err != nil {
			w.WriteError(errors.Wrap(err, "could not calculate total value").Error())

		}

		total := ethTotal + btcTotal + bchTotal
		aud, err := usdToAud(total)
		if err != nil {
			w.WriteError(errors.Wrap(err, "could not convert to AUD").Error())

		}
		w.Write(fmt.Sprintf("Your crypto net worth is: %.2f USD (%.2f AUD)", total, aud))
	}

}

func (cmd *PortfolioCommand) Run(cr *messagerouter.CommandRequest, rw messagerouter.ResponseWriter) error {
	msg := strings.Fields(cr.Text)

	err := cmd.App.Run(msg)
	if err != nil {
		return err
	}
	return nil

}

type CLIAdaptor interface {
	Initialise(cr *messagerouter.CommandRequest, rw messagerouter.ResponseWriter)
	Run(cr *messagerouter.CommandRequest, rw messagerouter.ResponseWriter) error
}

// PortfolioWrapper
func PortfolioWrapper(app CLIAdaptor) func(cr *messagerouter.CommandRequest, w messagerouter.ResponseWriter) error {
	return func(cr *messagerouter.CommandRequest, w messagerouter.ResponseWriter) error {
		app.Initialise(cr, w)
		app.Run(cr, w)
		return nil
	}
}

func calcValue(ownAmt, tickerAmt string) (float64, error) {
	own, err := strconv.ParseFloat(ownAmt, 64)
	if err != nil {
		return 0, err
	}
	tick, err := strconv.ParseFloat(tickerAmt, 64)
	if err != nil {
		return 0, err
	}
	return own * tick, err
}

type forexTicker struct {
	Base  string `json:"base"`
	Date  string `json:"date"`
	Rates struct {
		AUD float64 `json:"AUD"`
	} `json:"rates"`
}

func usdToAud(usd float64) (float64, error) {

	resp, err := http.Get("https://api.fixer.io/latest?symbols=USD,AUD")
	if err != nil {
		return 0, err
	}
	result := &forexTicker{}
	err = json.NewDecoder(resp.Body).Decode(result)
	if err != nil {
		return 0, err
	}

	converted := result.Rates.AUD * usd
	return converted, nil
}

func init() {
	Handle("!portfolio", PortfolioWrapper(&PortfolioCommand{}))
}
<|endoftext|>


<commit_before>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]struct{}
	status autonat.NATStatus
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]struct{}),
		disconnect: make(chan struct{}, 1),
		status:     autonat.NATStatusUnknown,
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	return ar.relayAddrs(ar.addrsF(addrs))
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			ar.mx.Lock()
			ar.status = autonat.NATStatusUnknown
			ar.mx.Unlock()
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			ar.mx.Lock()
			if ar.status != autonat.NATStatusPublic {
				push = true
			}
			ar.status = autonat.NATStatusPublic
			ar.mx.Unlock()

		case autonat.NATStatusPrivate:
			update := ar.findRelays(ctx)
			ar.mx.Lock()
			if update || ar.status != autonat.NATStatusPrivate {
				push = true
			}
			ar.status = autonat.NATStatusPrivate
			ar.mx.Unlock()
		}

		if push {
			push = false
			ar.host.PushIdentify()
		}

		select {
		case <-ar.disconnect:
			push = true
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) bool {
again:
	ar.mx.Lock()
	haveRelays := len(ar.relays)
	if haveRelays >= DesiredRelays {
		ar.mx.Unlock()
		return false
	}
	need := DesiredRelays - len(ar.relays)
	ar.mx.Unlock()

	limit := 1000

	dctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	pis, err := discovery.FindPeers(dctx, ar.discover, RelayRendezvous, limit)
	cancel()
	if err != nil {
		log.Debugf("error discovering relays: %s", err.Error())

		if haveRelays == 0 {
			log.Debug("no relays connected; retrying in 30s")
			select {
			case <-time.After(30 * time.Second):
				goto again
			case <-ctx.Done():
				return false
			}
		}
	}

	log.Debugf("discovered %d relays", len(pis))

	pis = ar.selectRelays(ctx, pis, 20, 50)
	update := 0

	for _, pi := range pis {
		ar.mx.Lock()
		if _, ok := ar.relays[pi.ID]; ok {
			ar.mx.Unlock()
			continue
		}
		ar.mx.Unlock()

		cctx, cancel := context.WithTimeout(ctx, 15*time.Second)
		err = ar.host.Connect(cctx, pi)
		cancel()
		if err != nil {
			log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
			continue
		}

		log.Debugf("connected to relay %s", pi.ID)
		ar.mx.Lock()
		ar.relays[pi.ID] = struct{}{}
		haveRelays++
		ar.mx.Unlock()

		// tag the connection as very important
		ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)

		update++
		need--
		if need == 0 {
			break
		}
	}

	if haveRelays == 0 {
		// we failed to find any relays and we are not connected to any!
		// wait a little and try again, the discovery query might have returned only dead peers
		log.Debug("no relays connected; retrying in 30s")
		select {
		case <-time.After(30 * time.Second):
			goto again
		case <-ctx.Done():
			return false
		}
	}

	return update > 0
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []pstore.PeerInfo, count, maxq int) []pstore.PeerInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	if len(pis) == 0 {
		return pis
	}

	if len(pis) < count {
		count = len(pis)
	}

	if len(pis) < maxq {
		maxq = len(pis)
	}

	// only select relays that can be found by routing
	type queryResult struct {
		pi  pstore.PeerInfo
		err error
	}
	result := make([]pstore.PeerInfo, 0, count)
	resultCh := make(chan queryResult, maxq)

	qctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	// shuffle to randomize the order of queries
	shuffleRelays(pis)
	for _, pi := range pis[:maxq] {
		// first check to see if we already know this peer from a previous query
		addrs := ar.host.Peerstore().Addrs(pi.ID)
		if len(addrs) > 0 {
			resultCh <- queryResult{pi: pstore.PeerInfo{ID: pi.ID, Addrs: addrs}, err: nil}
			continue
		}
		go func(p peer.ID) {
			pi, err := ar.router.FindPeer(qctx, p)
			if err != nil {
				log.Debugf("error finding relay peer %s: %s", p, err.Error())
			}
			resultCh <- queryResult{pi: pi, err: err}
		}(pi.ID)
	}

	rcount := 0
	for len(result) < count && rcount < maxq {
		select {
		case qr := <-resultCh:
			rcount++
			if qr.err == nil {
				result = append(result, qr.pi)
			}

		case <-qctx.Done():
			break
		}
	}

	shuffleRelays(result)
	return result
}

// This function is computes the NATed relay addrs when our status is private:
// - The public addrs are removed from the address set.
// - The non-public addrs are included verbatim so that peers behind the same NAT/firewall
//   can still dial us directly.
// - On top of those, we add the relay-specific addrs for the relays to which we are
//   connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
//   through which we can be dialed.
func (ar *AutoRelay) relayAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	if ar.status != autonat.NATStatusPrivate {
		ar.mx.Unlock()
		return addrs
	}

	relays := make([]peer.ID, 0, len(ar.relays))
	for p := range ar.relays {
		relays = append(relays, p)
	}
	ar.mx.Unlock()

	raddrs := make([]ma.Multiaddr, 0, 4*len(relays)+2)

	// only keep private addrs from the original addr set
	for _, addr := range addrs {
		if manet.IsPrivateAddr(addr) {
			raddrs = append(raddrs, addr)
		}
	}

	// add relay specific addrs to the list
	for _, p := range relays {
		addrs := cleanupAddressSet(ar.host.Peerstore().Addrs(p))

		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", p.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range addrs {
			pub := addr.Encapsulate(circuit)
			raddrs = append(raddrs, pub)
		}
	}

	return raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

// Notifee
func (ar *AutoRelay) Listen(inet.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(inet.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(inet.Network, inet.Conn)      {}

func (ar *AutoRelay) Disconnected(net inet.Network, c inet.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == inet.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(inet.Network, inet.Stream) {}
func (ar *AutoRelay) ClosedStream(inet.Network, inet.Stream) {}
<commit_msg>gate max number of retries in findRelays<commit_after>package relay

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	basic "github.com/libp2p/go-libp2p/p2p/host/basic"

	autonat "github.com/libp2p/go-libp2p-autonat"
	_ "github.com/libp2p/go-libp2p-circuit"
	discovery "github.com/libp2p/go-libp2p-discovery"
	inet "github.com/libp2p/go-libp2p-net"
	peer "github.com/libp2p/go-libp2p-peer"
	pstore "github.com/libp2p/go-libp2p-peerstore"
	routing "github.com/libp2p/go-libp2p-routing"
	ma "github.com/multiformats/go-multiaddr"
	manet "github.com/multiformats/go-multiaddr-net"
)

const (
	RelayRendezvous = "/libp2p/relay"
)

var (
	DesiredRelays = 3

	BootDelay = 20 * time.Second
)

// AutoRelay is a Host that uses relays for connectivity when a NAT is detected.
type AutoRelay struct {
	host     *basic.BasicHost
	discover discovery.Discoverer
	router   routing.PeerRouting
	autonat  autonat.AutoNAT
	addrsF   basic.AddrsFactory

	disconnect chan struct{}

	mx     sync.Mutex
	relays map[peer.ID]struct{}
	status autonat.NATStatus
}

func NewAutoRelay(ctx context.Context, bhost *basic.BasicHost, discover discovery.Discoverer, router routing.PeerRouting) *AutoRelay {
	ar := &AutoRelay{
		host:       bhost,
		discover:   discover,
		router:     router,
		addrsF:     bhost.AddrsFactory,
		relays:     make(map[peer.ID]struct{}),
		disconnect: make(chan struct{}, 1),
		status:     autonat.NATStatusUnknown,
	}
	ar.autonat = autonat.NewAutoNAT(ctx, bhost, ar.baseAddrs)
	bhost.AddrsFactory = ar.hostAddrs
	bhost.Network().Notify(ar)
	go ar.background(ctx)
	return ar
}

func (ar *AutoRelay) baseAddrs() []ma.Multiaddr {
	return ar.addrsF(ar.host.AllAddrs())
}

func (ar *AutoRelay) hostAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	return ar.relayAddrs(ar.addrsF(addrs))
}

func (ar *AutoRelay) background(ctx context.Context) {
	select {
	case <-time.After(autonat.AutoNATBootDelay + BootDelay):
	case <-ctx.Done():
		return
	}

	// when true, we need to identify push
	push := false

	for {
		wait := autonat.AutoNATRefreshInterval
		switch ar.autonat.Status() {
		case autonat.NATStatusUnknown:
			ar.mx.Lock()
			ar.status = autonat.NATStatusUnknown
			ar.mx.Unlock()
			wait = autonat.AutoNATRetryInterval

		case autonat.NATStatusPublic:
			ar.mx.Lock()
			if ar.status != autonat.NATStatusPublic {
				push = true
			}
			ar.status = autonat.NATStatusPublic
			ar.mx.Unlock()

		case autonat.NATStatusPrivate:
			update := ar.findRelays(ctx)
			ar.mx.Lock()
			if update || ar.status != autonat.NATStatusPrivate {
				push = true
			}
			ar.status = autonat.NATStatusPrivate
			ar.mx.Unlock()
		}

		if push {
			push = false
			ar.host.PushIdentify()
		}

		select {
		case <-ar.disconnect:
			push = true
		case <-time.After(wait):
		case <-ctx.Done():
			return
		}
	}
}

func (ar *AutoRelay) findRelays(ctx context.Context) bool {
	retry := 0

again:
	ar.mx.Lock()
	haveRelays := len(ar.relays)
	if haveRelays >= DesiredRelays {
		ar.mx.Unlock()
		return false
	}
	need := DesiredRelays - len(ar.relays)
	ar.mx.Unlock()

	limit := 1000

	dctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	pis, err := discovery.FindPeers(dctx, ar.discover, RelayRendezvous, limit)
	cancel()
	if err != nil {
		log.Debugf("error discovering relays: %s", err.Error())

		if haveRelays == 0 {
			retry++
			if retry > 5 {
				log.Debug("no relays connected; giving up")
				return false
			}

			log.Debug("no relays connected; retrying in 30s")
			select {
			case <-time.After(30 * time.Second):
				goto again
			case <-ctx.Done():
				return false
			}
		}
	}

	log.Debugf("discovered %d relays", len(pis))

	pis = ar.selectRelays(ctx, pis, 20, 50)
	update := 0

	for _, pi := range pis {
		ar.mx.Lock()
		if _, ok := ar.relays[pi.ID]; ok {
			ar.mx.Unlock()
			continue
		}
		ar.mx.Unlock()

		cctx, cancel := context.WithTimeout(ctx, 15*time.Second)
		err = ar.host.Connect(cctx, pi)
		cancel()
		if err != nil {
			log.Debugf("error connecting to relay %s: %s", pi.ID, err.Error())
			continue
		}

		log.Debugf("connected to relay %s", pi.ID)
		ar.mx.Lock()
		ar.relays[pi.ID] = struct{}{}
		haveRelays++
		ar.mx.Unlock()

		// tag the connection as very important
		ar.host.ConnManager().TagPeer(pi.ID, "relay", 42)

		update++
		need--
		if need == 0 {
			break
		}
	}

	if haveRelays == 0 {
		// we failed to find any relays and we are not connected to any!
		// wait a little and try again, the discovery query might have returned only dead peers
		retry++
		if retry > 5 {
			log.Debug("no relays connected; giving up")
			return false
		}

		log.Debug("no relays connected; retrying in 30s")
		select {
		case <-time.After(30 * time.Second):
			goto again
		case <-ctx.Done():
			return false
		}
	}

	return update > 0
}

func (ar *AutoRelay) selectRelays(ctx context.Context, pis []pstore.PeerInfo, count, maxq int) []pstore.PeerInfo {
	// TODO better relay selection strategy; this just selects random relays
	//      but we should probably use ping latency as the selection metric

	if len(pis) == 0 {
		return pis
	}

	if len(pis) < count {
		count = len(pis)
	}

	if len(pis) < maxq {
		maxq = len(pis)
	}

	// only select relays that can be found by routing
	type queryResult struct {
		pi  pstore.PeerInfo
		err error
	}
	result := make([]pstore.PeerInfo, 0, count)
	resultCh := make(chan queryResult, maxq)

	qctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	// shuffle to randomize the order of queries
	shuffleRelays(pis)
	for _, pi := range pis[:maxq] {
		// first check to see if we already know this peer from a previous query
		addrs := ar.host.Peerstore().Addrs(pi.ID)
		if len(addrs) > 0 {
			resultCh <- queryResult{pi: pstore.PeerInfo{ID: pi.ID, Addrs: addrs}, err: nil}
			continue
		}
		go func(p peer.ID) {
			pi, err := ar.router.FindPeer(qctx, p)
			if err != nil {
				log.Debugf("error finding relay peer %s: %s", p, err.Error())
			}
			resultCh <- queryResult{pi: pi, err: err}
		}(pi.ID)
	}

	rcount := 0
	for len(result) < count && rcount < maxq {
		select {
		case qr := <-resultCh:
			rcount++
			if qr.err == nil {
				result = append(result, qr.pi)
			}

		case <-qctx.Done():
			break
		}
	}

	shuffleRelays(result)
	return result
}

// This function is computes the NATed relay addrs when our status is private:
// - The public addrs are removed from the address set.
// - The non-public addrs are included verbatim so that peers behind the same NAT/firewall
//   can still dial us directly.
// - On top of those, we add the relay-specific addrs for the relays to which we are
//   connected. For each non-private relay addr, we encapsulate the p2p-circuit addr
//   through which we can be dialed.
func (ar *AutoRelay) relayAddrs(addrs []ma.Multiaddr) []ma.Multiaddr {
	ar.mx.Lock()
	if ar.status != autonat.NATStatusPrivate {
		ar.mx.Unlock()
		return addrs
	}

	relays := make([]peer.ID, 0, len(ar.relays))
	for p := range ar.relays {
		relays = append(relays, p)
	}
	ar.mx.Unlock()

	raddrs := make([]ma.Multiaddr, 0, 4*len(relays)+2)

	// only keep private addrs from the original addr set
	for _, addr := range addrs {
		if manet.IsPrivateAddr(addr) {
			raddrs = append(raddrs, addr)
		}
	}

	// add relay specific addrs to the list
	for _, p := range relays {
		addrs := cleanupAddressSet(ar.host.Peerstore().Addrs(p))

		circuit, err := ma.NewMultiaddr(fmt.Sprintf("/p2p/%s/p2p-circuit", p.Pretty()))
		if err != nil {
			panic(err)
		}

		for _, addr := range addrs {
			pub := addr.Encapsulate(circuit)
			raddrs = append(raddrs, pub)
		}
	}

	return raddrs
}

func shuffleRelays(pis []pstore.PeerInfo) {
	for i := range pis {
		j := rand.Intn(i + 1)
		pis[i], pis[j] = pis[j], pis[i]
	}
}

// Notifee
func (ar *AutoRelay) Listen(inet.Network, ma.Multiaddr)      {}
func (ar *AutoRelay) ListenClose(inet.Network, ma.Multiaddr) {}
func (ar *AutoRelay) Connected(inet.Network, inet.Conn)      {}

func (ar *AutoRelay) Disconnected(net inet.Network, c inet.Conn) {
	p := c.RemotePeer()

	ar.mx.Lock()
	defer ar.mx.Unlock()

	if ar.host.Network().Connectedness(p) == inet.Connected {
		// We have a second connection.
		return
	}

	if _, ok := ar.relays[p]; ok {
		delete(ar.relays, p)
		select {
		case ar.disconnect <- struct{}{}:
		default:
		}
	}
}

func (ar *AutoRelay) OpenedStream(inet.Network, inet.Stream) {}
func (ar *AutoRelay) ClosedStream(inet.Network, inet.Stream) {}
<|endoftext|>


<commit_before>package restapi

import (
	"net/http"
	"time"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/sebest/hooky/models"
	"gopkg.in/mgo.v2/bson"
)

// Application is a list of Tasks with a common application Name.
type Application struct {
	// ID is the ID of the Application.
	ID string `json:"id"`

	// Created is the date when the Application was created.
	Created string `json:"created"`

	// Account is the ID of the Account owning the Application.
	Account string `json:"account"`

	// Name is the application's name.
	Name string `json:"name"`
}

func applicationParams(r *rest.Request) (bson.ObjectId, string, error) {
	accountID, err := PathAccountID(r)
	if err != nil {
		return accountID, "", err
	}
	// TODO handle errors
	applicationName := r.PathParam("application")
	if applicationName == "" {
		applicationName = "default"
	}
	return accountID, applicationName, nil
}

// NewApplicationFromModel returns a Application object for use with the Rest API
// from a Application model.
func NewApplicationFromModel(application *models.Application) *Application {
	return &Application{
		ID:      application.ID.Hex(),
		Created: application.ID.Time().UTC().Format(time.RFC3339),
		Account: application.Account.Hex(),
		Name:    application.Name,
	}
}

// PutApplication ...
func PutApplication(w rest.ResponseWriter, r *rest.Request) {
	accountID, applicationName, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	rc := &Application{}
	if err := r.DecodeJsonPayload(rc); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	b := GetBase(r)
	application, err := b.NewApplication(accountID, applicationName)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.WriteJson(NewApplicationFromModel(application))
}

// GetApplication ...
func GetApplication(w rest.ResponseWriter, r *rest.Request) {
	accountID, applicationName, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := GetBase(r)
	application, err := b.GetApplication(accountID, applicationName)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if application == nil {
		rest.NotFound(w, r)
		return
	}
	w.WriteJson(NewApplicationFromModel(application))
}

// DeleteApplication ...
func DeleteApplication(w rest.ResponseWriter, r *rest.Request) {
	accountID, applicationName, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := GetBase(r)
	if err := b.DeleteApplication(accountID, applicationName); err != nil {
		if err == models.ErrDeleteDefaultApplication {
			rest.Error(w, err.Error(), http.StatusForbidden)
		} else {
			rest.Error(w, err.Error(), http.StatusInternalServerError)
		}
	}
}

// DeleteApplications ...
func DeleteApplications(w rest.ResponseWriter, r *rest.Request) {
	accountID, _, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := GetBase(r)
	if err := b.DeleteApplications(accountID); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// GetApplications ...
func GetApplications(w rest.ResponseWriter, r *rest.Request) {
	accountID, _, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := GetBase(r)
	lp := parseListQuery(r)
	var applications []*models.Application
	lr := &models.ListResult{
		List: &applications,
	}

	if err := b.GetApplications(accountID, lp, lr); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if lr.Count == 0 {
		rest.NotFound(w, r)
		return
	}
	rt := make([]*Application, len(applications))
	for idx, application := range applications {
		rt[idx] = NewApplicationFromModel(application)
	}
	w.WriteJson(models.ListResult{
		List:    rt,
		HasMore: lr.HasMore,
		Total:   lr.Total,
		Count:   lr.Count,
		Page:    lr.Page,
		Pages:   lr.Pages,
	})
}
<commit_msg>No JSON payload is not an error<commit_after>package restapi

import (
	"net/http"
	"time"

	"github.com/ant0ine/go-json-rest/rest"
	"github.com/sebest/hooky/models"
	"gopkg.in/mgo.v2/bson"
)

// Application is a list of Tasks with a common application Name.
type Application struct {
	// ID is the ID of the Application.
	ID string `json:"id"`

	// Created is the date when the Application was created.
	Created string `json:"created"`

	// Account is the ID of the Account owning the Application.
	Account string `json:"account"`

	// Name is the application's name.
	Name string `json:"name"`
}

func applicationParams(r *rest.Request) (bson.ObjectId, string, error) {
	accountID, err := PathAccountID(r)
	if err != nil {
		return accountID, "", err
	}
	// TODO handle errors
	applicationName := r.PathParam("application")
	if applicationName == "" {
		applicationName = "default"
	}
	return accountID, applicationName, nil
}

// NewApplicationFromModel returns a Application object for use with the Rest API
// from a Application model.
func NewApplicationFromModel(application *models.Application) *Application {
	return &Application{
		ID:      application.ID.Hex(),
		Created: application.ID.Time().UTC().Format(time.RFC3339),
		Account: application.Account.Hex(),
		Name:    application.Name,
	}
}

// PutApplication ...
func PutApplication(w rest.ResponseWriter, r *rest.Request) {
	accountID, applicationName, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	rc := &Application{}
	if err := r.DecodeJsonPayload(rc); err != nil {
		if err != rest.ErrJSONPayloadEmpty {
			rest.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
	b := GetBase(r)
	application, err := b.NewApplication(accountID, applicationName)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	w.WriteJson(NewApplicationFromModel(application))
}

// GetApplication ...
func GetApplication(w rest.ResponseWriter, r *rest.Request) {
	accountID, applicationName, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := GetBase(r)
	application, err := b.GetApplication(accountID, applicationName)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if application == nil {
		rest.NotFound(w, r)
		return
	}
	w.WriteJson(NewApplicationFromModel(application))
}

// DeleteApplication ...
func DeleteApplication(w rest.ResponseWriter, r *rest.Request) {
	accountID, applicationName, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := GetBase(r)
	if err := b.DeleteApplication(accountID, applicationName); err != nil {
		if err == models.ErrDeleteDefaultApplication {
			rest.Error(w, err.Error(), http.StatusForbidden)
		} else {
			rest.Error(w, err.Error(), http.StatusInternalServerError)
		}
	}
}

// DeleteApplications ...
func DeleteApplications(w rest.ResponseWriter, r *rest.Request) {
	accountID, _, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := GetBase(r)
	if err := b.DeleteApplications(accountID); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

// GetApplications ...
func GetApplications(w rest.ResponseWriter, r *rest.Request) {
	accountID, _, err := applicationParams(r)
	if err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	b := GetBase(r)
	lp := parseListQuery(r)
	var applications []*models.Application
	lr := &models.ListResult{
		List: &applications,
	}

	if err := b.GetApplications(accountID, lp, lr); err != nil {
		rest.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	if lr.Count == 0 {
		rest.NotFound(w, r)
		return
	}
	rt := make([]*Application, len(applications))
	for idx, application := range applications {
		rt[idx] = NewApplicationFromModel(application)
	}
	w.WriteJson(models.ListResult{
		List:    rt,
		HasMore: lr.HasMore,
		Total:   lr.Total,
		Count:   lr.Count,
		Page:    lr.Page,
		Pages:   lr.Pages,
	})
}
<|endoftext|>


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

package govppmux

import (
	"context"
	"errors"
	"sync"
	"time"

	"git.fd.io/govpp.git/adapter"
	"git.fd.io/govpp.git/api"
	govpp "git.fd.io/govpp.git/core"
	"github.com/ligato/cn-infra/flavors/local"
	"github.com/ligato/cn-infra/health/statuscheck"
	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/logging/logrus"
	"github.com/ligato/vpp-agent/plugins/govppmux/vppcalls"
)

// GOVPPPlugin implements the govppmux plugin interface.
type GOVPPPlugin struct {
	Deps // Inject.

	vppConn    *govpp.Connection
	vppAdapter adapter.VppAdapter
	vppConChan chan govpp.ConnectionEvent

	vppAPIChan *api.Channel

	cancel context.CancelFunc // Cancel can be used to cancel all goroutines and their jobs inside of the plugin.
	wg     sync.WaitGroup     // Wait group allows to wait until all goroutines of the plugin have finished.
}

// Deps groups injected dependencies of plugin
// so that they do not mix with other plugin fields.
type Deps struct {
	local.PluginInfraDeps // inject
}

// Config groups the configurable parameter of GoVpp.
type Config struct {
	HealthCheckProbeInterval time.Duration `json:"health-check-probe-interval"`
	HealthCheckReplyTimeout  time.Duration `json:"health-check-reply-timeout"`
	HealthCheckThreshold     int           `json:"health-check-threshold"`
}

// FromExistingAdapter is used mainly for testing purposes.
func FromExistingAdapter(vppAdapter adapter.VppAdapter) *GOVPPPlugin {
	ret := &GOVPPPlugin{
		vppAdapter: vppAdapter,
	}
	return ret
}

// Init is the entry point called by Agent Core. A single binary-API connection to VPP is established.
func (plugin *GOVPPPlugin) Init() error {
	var err error

	govppLogger := plugin.Deps.Log.NewLogger("GoVpp")
	if govppLogger, ok := govppLogger.(*logrus.Logger); ok {
		govppLogger.SetLevel(logging.InfoLevel)
		govpp.SetLogger(govppLogger.StandardLogger())
	}

	plugin.PluginName = plugin.Deps.PluginName

	cfg := defaultConfig()
	found, err := plugin.PluginConfig.GetValue(&cfg)
	if err != nil {
		return err
	}
	if found {
		govpp.SetHealthCheckProbeInterval(cfg.HealthCheckProbeInterval)
		govpp.SetHealthCheckReplyTimeout(cfg.HealthCheckReplyTimeout)
		govpp.SetHealthCheckThreshold(cfg.HealthCheckThreshold)
		plugin.Log.Debug("Setting govpp parameters", cfg)
	}

	if plugin.vppAdapter == nil {
		plugin.vppAdapter = NewVppAdapter()
	} else {
		plugin.Log.Info("Reusing existing vppAdapter") //this is used for testing purposes
	}

	startTime := time.Now()
	plugin.vppConn, plugin.vppConChan, err = govpp.AsyncConnect(plugin.vppAdapter)
	if err != nil {
		return err
	}

	// TODO: Async connect & automatic reconnect support is not yet implemented in the agent,
	// so synchronously wait until connected to VPP.
	status := <-plugin.vppConChan
	if status.State != govpp.Connected {
		return errors.New("unable to connect to VPP")
	}
	vppConnectTime := time.Since(startTime)
	plugin.Log.WithField("durationInNs", vppConnectTime.Nanoseconds()).Info("Connecting to VPP took ", vppConnectTime)
	plugin.retrieveVersion()

	// Register providing status reports (push mode)
	plugin.StatusCheck.Register(plugin.PluginName, nil)
	plugin.StatusCheck.ReportStateChange(plugin.PluginName, statuscheck.OK, nil)
	plugin.Log.Debug("govpp connect success ", plugin.vppConn)

	var ctx context.Context
	ctx, plugin.cancel = context.WithCancel(context.Background())
	go plugin.handleVPPConnectionEvents(ctx)

	return nil
}

// Close cleans up the resources allocated by the govppmux plugin.
func (plugin *GOVPPPlugin) Close() error {
	plugin.cancel()
	plugin.wg.Wait()

	defer func() {
		if plugin.vppConn != nil {
			plugin.vppConn.Disconnect()
		}
	}()

	return nil
}

// NewAPIChannel returns a new API channel for communication with VPP via govpp core.
// It uses default buffer sizes for the request and reply Go channels.
//
// Example of binary API call from some plugin using GOVPP:
//      ch, _ := govpp_mux.NewAPIChannel()
//      ch.SendRequest(req).ReceiveReply
func (plugin *GOVPPPlugin) NewAPIChannel() (*api.Channel, error) {
	return plugin.vppConn.NewAPIChannel()
}

// NewAPIChannelBuffered returns a new API channel for communication with VPP via govpp core.
// It allows to specify custom buffer sizes for the request and reply Go channels.
//
// Example of binary API call from some plugin using GOVPP:
//      ch, _ := govpp_mux.NewAPIChannelBuffered(100, 100)
//      ch.SendRequest(req).ReceiveReply
func (plugin *GOVPPPlugin) NewAPIChannelBuffered(reqChanBufSize, replyChanBufSize int) (*api.Channel, error) {
	return plugin.vppConn.NewAPIChannelBuffered(reqChanBufSize, replyChanBufSize)
}

// handleVPPConnectionEvents handles VPP connection events.
func (plugin *GOVPPPlugin) handleVPPConnectionEvents(ctx context.Context) {
	plugin.wg.Add(1)
	defer plugin.wg.Done()

	// TODO: support for VPP reconnect

	for {
		select {
		case status := <-plugin.vppConChan:
			if status.State == govpp.Connected {
				plugin.retrieveVersion()
				plugin.StatusCheck.ReportStateChange(plugin.PluginName, statuscheck.OK, nil)
			} else {
				plugin.StatusCheck.ReportStateChange(plugin.PluginName, statuscheck.Error, errors.New("VPP disconnected"))
			}

		case <-ctx.Done():
			return
		}
	}
}

func (plugin *GOVPPPlugin) retrieveVersion() {
	if plugin.vppAPIChan == nil {
		var err error
		if plugin.vppAPIChan, err = plugin.vppConn.NewAPIChannel(); err != nil {
			plugin.Log.Error("getting new api channel failed:", err)
			return
		}
	}

	info, err := vppcalls.GetVersionInfo(plugin.Log, plugin.vppAPIChan)
	if err != nil {
		plugin.Log.Warn("getting version info failed:", err)
		return
	}

	plugin.Log.Debugf("version info: %+v", info)
	plugin.Log.Infof("VPP version: %v (%v)", info.Version, info.BuildDate)
}

func defaultConfig() Config {
	c := Config{
		HealthCheckProbeInterval: time.Second,
		HealthCheckReplyTimeout:  100 * time.Millisecond,
		HealthCheckThreshold:     1,
	}
	return c
}
<commit_msg>Call close on VPP API channel<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package govppmux

import (
	"context"
	"errors"
	"sync"
	"time"

	"git.fd.io/govpp.git/adapter"
	"git.fd.io/govpp.git/api"
	govpp "git.fd.io/govpp.git/core"
	"github.com/ligato/cn-infra/flavors/local"
	"github.com/ligato/cn-infra/health/statuscheck"
	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/logging/logrus"
	"github.com/ligato/vpp-agent/plugins/govppmux/vppcalls"
)

// GOVPPPlugin implements the govppmux plugin interface.
type GOVPPPlugin struct {
	Deps // Inject.

	vppConn    *govpp.Connection
	vppAdapter adapter.VppAdapter
	vppConChan chan govpp.ConnectionEvent

	cancel context.CancelFunc // Cancel can be used to cancel all goroutines and their jobs inside of the plugin.
	wg     sync.WaitGroup     // Wait group allows to wait until all goroutines of the plugin have finished.
}

// Deps groups injected dependencies of plugin
// so that they do not mix with other plugin fields.
type Deps struct {
	local.PluginInfraDeps // inject
}

// Config groups the configurable parameter of GoVpp.
type Config struct {
	HealthCheckProbeInterval time.Duration `json:"health-check-probe-interval"`
	HealthCheckReplyTimeout  time.Duration `json:"health-check-reply-timeout"`
	HealthCheckThreshold     int           `json:"health-check-threshold"`
}

// FromExistingAdapter is used mainly for testing purposes.
func FromExistingAdapter(vppAdapter adapter.VppAdapter) *GOVPPPlugin {
	ret := &GOVPPPlugin{
		vppAdapter: vppAdapter,
	}
	return ret
}

// Init is the entry point called by Agent Core. A single binary-API connection to VPP is established.
func (plugin *GOVPPPlugin) Init() error {
	var err error

	govppLogger := plugin.Deps.Log.NewLogger("GoVpp")
	if govppLogger, ok := govppLogger.(*logrus.Logger); ok {
		govppLogger.SetLevel(logging.InfoLevel)
		govpp.SetLogger(govppLogger.StandardLogger())
	}

	plugin.PluginName = plugin.Deps.PluginName

	cfg := defaultConfig()
	found, err := plugin.PluginConfig.GetValue(&cfg)
	if err != nil {
		return err
	}
	if found {
		govpp.SetHealthCheckProbeInterval(cfg.HealthCheckProbeInterval)
		govpp.SetHealthCheckReplyTimeout(cfg.HealthCheckReplyTimeout)
		govpp.SetHealthCheckThreshold(cfg.HealthCheckThreshold)
		plugin.Log.Debug("Setting govpp parameters", cfg)
	}

	if plugin.vppAdapter == nil {
		plugin.vppAdapter = NewVppAdapter()
	} else {
		plugin.Log.Info("Reusing existing vppAdapter") //this is used for testing purposes
	}

	startTime := time.Now()
	plugin.vppConn, plugin.vppConChan, err = govpp.AsyncConnect(plugin.vppAdapter)
	if err != nil {
		return err
	}

	// TODO: Async connect & automatic reconnect support is not yet implemented in the agent,
	// so synchronously wait until connected to VPP.
	status := <-plugin.vppConChan
	if status.State != govpp.Connected {
		return errors.New("unable to connect to VPP")
	}
	vppConnectTime := time.Since(startTime)
	plugin.Log.WithField("durationInNs", vppConnectTime.Nanoseconds()).Info("Connecting to VPP took ", vppConnectTime)
	plugin.retrieveVersion()

	// Register providing status reports (push mode)
	plugin.StatusCheck.Register(plugin.PluginName, nil)
	plugin.StatusCheck.ReportStateChange(plugin.PluginName, statuscheck.OK, nil)
	plugin.Log.Debug("govpp connect success ", plugin.vppConn)

	var ctx context.Context
	ctx, plugin.cancel = context.WithCancel(context.Background())
	go plugin.handleVPPConnectionEvents(ctx)

	return nil
}

// Close cleans up the resources allocated by the govppmux plugin.
func (plugin *GOVPPPlugin) Close() error {
	plugin.cancel()
	plugin.wg.Wait()

	defer func() {
		if plugin.vppConn != nil {
			plugin.vppConn.Disconnect()
		}
	}()

	return nil
}

// NewAPIChannel returns a new API channel for communication with VPP via govpp core.
// It uses default buffer sizes for the request and reply Go channels.
//
// Example of binary API call from some plugin using GOVPP:
//      ch, _ := govpp_mux.NewAPIChannel()
//      ch.SendRequest(req).ReceiveReply
func (plugin *GOVPPPlugin) NewAPIChannel() (*api.Channel, error) {
	return plugin.vppConn.NewAPIChannel()
}

// NewAPIChannelBuffered returns a new API channel for communication with VPP via govpp core.
// It allows to specify custom buffer sizes for the request and reply Go channels.
//
// Example of binary API call from some plugin using GOVPP:
//      ch, _ := govpp_mux.NewAPIChannelBuffered(100, 100)
//      ch.SendRequest(req).ReceiveReply
func (plugin *GOVPPPlugin) NewAPIChannelBuffered(reqChanBufSize, replyChanBufSize int) (*api.Channel, error) {
	return plugin.vppConn.NewAPIChannelBuffered(reqChanBufSize, replyChanBufSize)
}

// handleVPPConnectionEvents handles VPP connection events.
func (plugin *GOVPPPlugin) handleVPPConnectionEvents(ctx context.Context) {
	plugin.wg.Add(1)
	defer plugin.wg.Done()

	// TODO: support for VPP reconnect

	for {
		select {
		case status := <-plugin.vppConChan:
			if status.State == govpp.Connected {
				plugin.retrieveVersion()
				plugin.StatusCheck.ReportStateChange(plugin.PluginName, statuscheck.OK, nil)
			} else {
				plugin.StatusCheck.ReportStateChange(plugin.PluginName, statuscheck.Error, errors.New("VPP disconnected"))
			}

		case <-ctx.Done():
			return
		}
	}
}

func (plugin *GOVPPPlugin) retrieveVersion() {
	vppAPIChan, err := plugin.vppConn.NewAPIChannel()
	if err != nil {
		plugin.Log.Error("getting new api channel failed:", err)
		return
	}
	defer vppAPIChan.Close()

	info, err := vppcalls.GetVersionInfo(plugin.Log, vppAPIChan)
	if err != nil {
		plugin.Log.Warn("getting version info failed:", err)
		return
	}

	plugin.Log.Debugf("version info: %+v", info)
	plugin.Log.Infof("VPP version: %v (%v)", info.Version, info.BuildDate)
}

func defaultConfig() Config {
	c := Config{
		HealthCheckProbeInterval: time.Second,
		HealthCheckReplyTimeout:  100 * time.Millisecond,
		HealthCheckThreshold:     1,
	}
	return c
}
<|endoftext|>


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

// +build darwin freebsd linux netbsd openbsd

package signal

import (
	"os"
	"runtime"
	"syscall"
	"testing"
	"time"
)

const sighup = syscall.SIGHUP

func waitSig(t *testing.T, c <-chan os.Signal, sig os.Signal) {
	select {
	case s := <-c:
		if s != sig {
			t.Fatalf("signal was %v, want %v", s, sig)
		}
	case <-time.After(1 * time.Second):
		t.Fatalf("timeout waiting for %v", sig)
	}
}

func TestSignal(t *testing.T) {
	// Ask for SIGHUP
	c := make(chan os.Signal, 1)
	Notify(c, sighup)

	t.Logf("sighup...")
	// Send this process a SIGHUP
	syscall.Kill(syscall.Getpid(), sighup)
	waitSig(t, c, sighup)

	// Ask for everything we can get.
	c1 := make(chan os.Signal, 1)
	Notify(c1)

	t.Logf("sigwinch...")
	// Send this process a SIGWINCH
	syscall.Kill(syscall.Getpid(), syscall.SIGWINCH)
	waitSig(t, c1, syscall.SIGWINCH)

	// Send two more SIGHUPs, to make sure that
	// they get delivered on c1 and that not reading
	// from c does not block everything.
	t.Logf("sigwinch...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c1, syscall.SIGHUP)
	t.Logf("sigwinch...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c1, syscall.SIGHUP)

	// The first SIGHUP should be waiting for us on c.
	waitSig(t, c, syscall.SIGHUP)
}

func TestStress(t *testing.T) {
	dur := 3 * time.Second
	if testing.Short() {
		dur = 100 * time.Millisecond
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(4))
	done := make(chan bool)
	finished := make(chan bool)
	go func() {
		sig := make(chan os.Signal, 1)
		Notify(sig, syscall.SIGUSR1)
	Loop:
		for {
			select {
			case <-sig:
			case <-done:
				break Loop
			}
		}
		finished <- true
	}()
	go func() {
	Loop:
		for {
			select {
			case <-done:
				break Loop
			default:
				syscall.Kill(syscall.Getpid(), syscall.SIGUSR1)
				runtime.Gosched()
			}
		}
		finished <- true
	}()
	time.Sleep(dur)
	close(done)
	<-finished
	<-finished
}
<commit_msg>os/signal: deflake test Fixes issue 4987.<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build darwin freebsd linux netbsd openbsd

package signal

import (
	"os"
	"runtime"
	"syscall"
	"testing"
	"time"
)

const sighup = syscall.SIGHUP

func waitSig(t *testing.T, c <-chan os.Signal, sig os.Signal) {
	select {
	case s := <-c:
		if s != sig {
			t.Fatalf("signal was %v, want %v", s, sig)
		}
	case <-time.After(1 * time.Second):
		t.Fatalf("timeout waiting for %v", sig)
	}
}

func TestSignal(t *testing.T) {
	// Ask for SIGHUP
	c := make(chan os.Signal, 1)
	Notify(c, sighup)

	t.Logf("sighup...")
	// Send this process a SIGHUP
	syscall.Kill(syscall.Getpid(), sighup)
	waitSig(t, c, sighup)

	// Ask for everything we can get.
	c1 := make(chan os.Signal, 1)
	Notify(c1)

	t.Logf("sigwinch...")
	// Send this process a SIGWINCH
	syscall.Kill(syscall.Getpid(), syscall.SIGWINCH)
	waitSig(t, c1, syscall.SIGWINCH)

	// Send two more SIGHUPs, to make sure that
	// they get delivered on c1 and that not reading
	// from c does not block everything.
	t.Logf("sigwinch...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c1, syscall.SIGHUP)
	t.Logf("sigwinch...")
	syscall.Kill(syscall.Getpid(), syscall.SIGHUP)
	waitSig(t, c1, syscall.SIGHUP)

	// The first SIGHUP should be waiting for us on c.
	waitSig(t, c, syscall.SIGHUP)
}

func TestStress(t *testing.T) {
	dur := 3 * time.Second
	if testing.Short() {
		dur = 100 * time.Millisecond
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(4))
	done := make(chan bool)
	finished := make(chan bool)
	go func() {
		sig := make(chan os.Signal, 1)
		Notify(sig, syscall.SIGUSR1)
	Loop:
		for {
			select {
			case <-sig:
			case <-done:
				break Loop
			}
		}
		finished <- true
	}()
	go func() {
	Loop:
		for {
			select {
			case <-done:
				break Loop
			default:
				syscall.Kill(syscall.Getpid(), syscall.SIGUSR1)
				runtime.Gosched()
			}
		}
		finished <- true
	}()
	time.Sleep(dur)
	close(done)
	<-finished
	<-finished
	// When run with 'go test -cpu=1,2,4' SIGUSR1 from this test can slip
	// into subsequent TestSignal() causing failure.
	// Sleep for a while to reduce the possibility of the failure.
	time.Sleep(10 * time.Millisecond)
}
<|endoftext|>


<commit_before>package rockredis

import (
	"errors"
	"fmt"
	"github.com/absolute8511/ZanRedisDB/common"
	"github.com/absolute8511/gorocksdb"
	"io"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

func GetCheckpointDir(term uint64, index uint64) string {
	return fmt.Sprintf("%016x-%016x", term, index)
}

type RockConfig struct {
	DataDir          string
	DefaultReadOpts  *gorocksdb.ReadOptions
	DefaultWriteOpts *gorocksdb.WriteOptions
}

func NewRockConfig() *RockConfig {
	c := &RockConfig{
		DefaultReadOpts:  gorocksdb.NewDefaultReadOptions(),
		DefaultWriteOpts: gorocksdb.NewDefaultWriteOptions(),
	}
	c.DefaultReadOpts.SetVerifyChecksums(false)
	return c
}

type RockDB struct {
	cfg              *RockConfig
	eng              *gorocksdb.DB
	dbOpts           *gorocksdb.Options
	defaultWriteOpts *gorocksdb.WriteOptions
	defaultReadOpts  *gorocksdb.ReadOptions
	wb               *gorocksdb.WriteBatch
	quit             chan struct{}
	wg               sync.WaitGroup
	backupC          chan *BackupInfo
}

func OpenRockDB(cfg *RockConfig) (*RockDB, error) {
	if len(cfg.DataDir) == 0 {
		return nil, errors.New("config error")
	}

	os.MkdirAll(cfg.DataDir, common.DIR_PERM)
	// options need be adjust due to using hdd or sdd, please reference
	// https://github.com/facebook/rocksdb/wiki/RocksDB-Tuning-Guide
	bbto := gorocksdb.NewDefaultBlockBasedTableOptions()
	// use large block to reduce index block size for hdd
	// if using ssd, should use the default value
	bbto.SetBlockSize(1024 * 16)
	// should about 20% less than host RAM
	// http://smalldatum.blogspot.com/2016/09/tuning-rocksdb-block-cache.html
	bbto.SetBlockCache(gorocksdb.NewLRUCache(1024 * 1024 * 1024))
	// for hdd , we nee cache index and filter blocks
	bbto.SetCacheIndexAndFilterBlocks(true)
	filter := gorocksdb.NewBloomFilter(10)
	bbto.SetFilterPolicy(filter)
	opts := gorocksdb.NewDefaultOptions()
	opts.SetBlockBasedTableFactory(bbto)
	opts.SetCreateIfMissing(true)
	opts.SetMaxOpenFiles(-1)
	// keep level0_file_num_compaction_trigger * write_buffer_size = max_bytes_for_level_base to minimize write amplification
	opts.SetWriteBufferSize(1024 * 1024 * 128)
	opts.SetMaxWriteBufferNumber(8)
	opts.SetLevel0FileNumCompactionTrigger(4)
	opts.SetMaxBytesForLevelBase(1024 * 1024 * 1024 * 2)
	opts.SetMinWriteBufferNumberToMerge(2)
	opts.SetTargetFileSizeBase(1024 * 1024 * 128)
	opts.SetMaxBackgroundFlushes(2)
	opts.SetMaxBackgroundCompactions(4)
	opts.SetMinLevelToCompress(3)
	// we use table, so we use prefix seek feature
	opts.SetPrefixExtractor(gorocksdb.NewFixedPrefixTransform(3))
	opts.SetMemtablePrefixBloomSizeRatio(0.1)
	opts.EnableStatistics()
	opts.SetMaxLogFileSize(1024 * 1024 * 32)
	opts.SetLogFileTimeToRoll(3600 * 24 * 3)
	// https://github.com/facebook/mysql-5.6/wiki/my.cnf-tuning
	// rate limiter need to reduce the compaction io

	db := &RockDB{
		cfg:              cfg,
		dbOpts:           opts,
		defaultReadOpts:  cfg.DefaultReadOpts,
		defaultWriteOpts: cfg.DefaultWriteOpts,
		wb:               gorocksdb.NewWriteBatch(),
		backupC:          make(chan *BackupInfo),
		quit:             make(chan struct{}),
	}
	eng, err := gorocksdb.OpenDb(opts, db.GetDataDir())
	if err != nil {
		return nil, err
	}
	db.eng = eng
	os.MkdirAll(db.GetBackupDir(), common.DIR_PERM)

	db.wg.Add(1)
	go func() {
		defer db.wg.Done()
		db.backupLoop()
	}()
	return db, nil
}

func GetBackupDir(base string) string {
	return path.Join(base, "rocksdb_backup")
}

func (r *RockDB) GetBackupBase() string {
	return r.cfg.DataDir
}

func (r *RockDB) GetBackupDir() string {
	return GetBackupDir(r.cfg.DataDir)
}

func (r *RockDB) GetDataDir() string {
	return path.Join(r.cfg.DataDir, "rocksdb")
}

func (r *RockDB) reOpen() error {
	var err error
	r.eng, err = gorocksdb.OpenDb(r.dbOpts, r.GetDataDir())
	return err
}

func (r *RockDB) CompactRange() {
	var rg gorocksdb.Range
	r.eng.CompactRange(rg)
}

func (r *RockDB) Close() {
	close(r.quit)
	r.wg.Wait()
	if r.defaultReadOpts != nil {
		r.defaultReadOpts.Destroy()
	}
	if r.defaultWriteOpts != nil {
		r.defaultWriteOpts.Destroy()
	}
	if r.eng != nil {
		r.eng.Close()
	}
}

func (r *RockDB) SetPerfLevel(level int) {
	// TODO:
}

func (r *RockDB) GetStatistics() string {
	return r.dbOpts.GetStatistics()
}

func (r *RockDB) GetInternalStatus() map[string]interface{} {
	status := make(map[string]interface{})
	bbt := r.dbOpts.GetBlockBasedTableFactory()
	if bbt != nil {
		bc := bbt.GetBlockCache()
		if bc != nil {
			status["block-cache-usage"] = bc.GetUsage()
			status["block-cache-pinned-usage"] = bc.GetPinnedUsage()
		}
	}

	memStr := r.eng.GetProperty("rocksdb.estimate-table-readers-mem")
	status["estimate-table-readers-mem"] = memStr
	memStr = r.eng.GetProperty("rocksdb.cur-size-all-mem-tables")
	status["cur-size-all-mem-tables"] = memStr
	memStr = r.eng.GetProperty("rocksdb.cur-size-active-mem-table")
	status["cur-size-active-mem-tables"] = memStr
	return status
}

func (r *RockDB) GetInternalPropertyStatus(p string) string {
	return r.eng.GetProperty(p)
}

func (r *RockDB) ReadRange(sKey, eKey []byte, maxNum int) chan common.KVRecord {
	retChan := make(chan common.KVRecord, 32)
	go func() {
		it := NewDBRangeLimitIterator(r.eng, sKey, eKey, RangeClose, 0, maxNum, false)
		defer it.Close()
		for it = it; it.Valid(); it.Next() {
			key := it.Key()
			value := it.Value()
			retChan <- common.KVRecord{Key: key, Value: value}
		}
		close(retChan)
	}()
	return retChan
}

type BackupInfo struct {
	backupDir string
	started   chan struct{}
	done      chan struct{}
	rsp       []byte
	err       error
}

func newBackupInfo(dir string) *BackupInfo {
	return &BackupInfo{
		backupDir: dir,
		started:   make(chan struct{}),
		done:      make(chan struct{}),
	}
}

func (self *BackupInfo) WaitReady() {
	select {
	case <-self.started:
	case <-self.done:
	}
}

func (self *BackupInfo) GetResult() ([]byte, error) {
	select {
	case <-self.done:
	}
	return self.rsp, self.err
}

func (r *RockDB) backupLoop() {
	for {
		select {
		case rsp, ok := <-r.backupC:
			if !ok {
				return
			}

			func() {
				defer close(rsp.done)
				log.Printf("begin backup to:%v \n", rsp.backupDir)
				start := time.Now()
				ck, err := gorocksdb.NewCheckpoint(r.eng)
				if err != nil {
					log.Printf("init checkpoint failed: %v", err)
					rsp.err = err
					return
				}
				_, err = os.Stat(rsp.backupDir)
				if !os.IsNotExist(err) {
					log.Printf("checkpoint exist: %v, remove it", rsp.backupDir)
					os.RemoveAll(rsp.backupDir)
				}
				time.AfterFunc(time.Second*2, func() {
					close(rsp.started)
				})
				err = ck.Save(rsp.backupDir)
				if err != nil {
					log.Printf("save checkpoint failed: %v", err)
					rsp.err = err
					return
				}
				cost := time.Now().Sub(start)
				log.Printf("backup done (cost %v), check point to: %v\n", cost.String(), rsp.backupDir)
				// TODO: purge some old checkpoint
				rsp.rsp = []byte(rsp.backupDir)
			}()
		case <-r.quit:
			return
		}
	}
}

func (r *RockDB) Backup(term uint64, index uint64) *BackupInfo {
	fname := GetCheckpointDir(term, index)
	checkpointDir := path.Join(r.GetBackupDir(), fname)
	bi := newBackupInfo(checkpointDir)
	select {
	case r.backupC <- bi:
	default:
		return nil
	}
	return bi
}

func (r *RockDB) IsLocalBackupOK(term uint64, index uint64) (bool, error) {
	backupDir := r.GetBackupDir()
	checkpointDir := GetCheckpointDir(term, index)
	ro := *r.dbOpts
	ro.SetCreateIfMissing(false)
	db, err := gorocksdb.OpenDbForReadOnly(&ro, path.Join(backupDir, checkpointDir), false)
	if err != nil {
		log.Printf("checkpoint open failed: %v", err)
		return false, err
	}
	db.Close()
	return true, nil
}

func copyFile(src, dst string, override bool) error {
	sfi, err := os.Stat(src)
	if err != nil {
		return err
	}
	if !sfi.Mode().IsRegular() {
		return fmt.Errorf("copyfile: non-regular source file %v (%v)", sfi.Name(), sfi.Mode().String())
	}
	_, err = os.Stat(dst)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	} else {
		if !override {
			return nil
		}
	}
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	_, err = io.Copy(out, in)
	if err != nil {
		out.Close()
		return err
	}
	err = out.Sync()
	if err != nil {
		out.Close()
		return err
	}
	return out.Close()
}

func (r *RockDB) Restore(term uint64, index uint64) error {
	// TODO: maybe write meta (snap term and index) and check the meta data in the backup
	backupDir := r.GetBackupDir()
	hasBackup, _ := r.IsLocalBackupOK(term, index)
	if !hasBackup {
		return errors.New("no backup for restore")
	}

	checkpointDir := GetCheckpointDir(term, index)
	start := time.Now()
	log.Printf("begin restore from checkpoint: %v\n", checkpointDir)
	r.eng.Close()
	// 1. remove all files in current db except sst files
	// 2. get the list of sst in checkpoint
	// 3. remove all the sst files not in the checkpoint list
	// 4. copy all files from checkpoint to current db and do not override sst
	matchName := path.Join(r.GetDataDir(), "*")
	nameList, err := filepath.Glob(matchName)
	if err != nil {
		log.Printf("list files failed:  %v\n", err)
		return err
	}
	ckNameList, err := filepath.Glob(path.Join(backupDir, checkpointDir, "*"))
	if err != nil {
		log.Printf("list checkpoint files failed:  %v\n", err)
		return err
	}
	ckSstNameMap := make(map[string]bool)
	for _, fn := range ckNameList {
		if strings.HasSuffix(fn, ".sst") {
			ckSstNameMap[fn] = true
		}
	}

	for _, fn := range nameList {
		shortName := path.Base(fn)
		if strings.HasPrefix(shortName, "LOG") {
			continue
		}
		if strings.HasSuffix(fn, ".sst") {
			if _, ok := ckSstNameMap[fn]; ok {
				log.Printf("keeping sst file: %v", fn)
				continue
			}
		}
		log.Printf("removing: %v", fn)
		os.RemoveAll(fn)
	}
	for _, fn := range ckNameList {
		dst := path.Join(r.GetDataDir(), path.Base(fn))
		err := copyFile(fn, dst, false)
		if err != nil {
			log.Printf("copy %v to %v failed: %v", fn, dst, err)
			return err
		} else {
			log.Printf("copy %v to %v done", fn, dst)
		}
	}

	err = r.reOpen()
	log.Printf("restore done, cost: %v\n", time.Now().Sub(start))
	if err != nil {
		log.Printf("reopen the restored db failed:  %v\n", err)
	}
	return err
}

func (r *RockDB) ClearBackup(term uint64, index uint64) error {
	backupDir := r.GetBackupDir()
	checkpointDir := GetCheckpointDir(term, index)
	return os.RemoveAll(path.Join(backupDir, checkpointDir))
}
<commit_msg>fix restore keep sst files<commit_after>package rockredis

import (
	"errors"
	"fmt"
	"github.com/absolute8511/ZanRedisDB/common"
	"github.com/absolute8511/gorocksdb"
	"io"
	"log"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

func GetCheckpointDir(term uint64, index uint64) string {
	return fmt.Sprintf("%016x-%016x", term, index)
}

type RockConfig struct {
	DataDir          string
	DefaultReadOpts  *gorocksdb.ReadOptions
	DefaultWriteOpts *gorocksdb.WriteOptions
}

func NewRockConfig() *RockConfig {
	c := &RockConfig{
		DefaultReadOpts:  gorocksdb.NewDefaultReadOptions(),
		DefaultWriteOpts: gorocksdb.NewDefaultWriteOptions(),
	}
	c.DefaultReadOpts.SetVerifyChecksums(false)
	return c
}

type RockDB struct {
	cfg              *RockConfig
	eng              *gorocksdb.DB
	dbOpts           *gorocksdb.Options
	defaultWriteOpts *gorocksdb.WriteOptions
	defaultReadOpts  *gorocksdb.ReadOptions
	wb               *gorocksdb.WriteBatch
	quit             chan struct{}
	wg               sync.WaitGroup
	backupC          chan *BackupInfo
}

func OpenRockDB(cfg *RockConfig) (*RockDB, error) {
	if len(cfg.DataDir) == 0 {
		return nil, errors.New("config error")
	}

	os.MkdirAll(cfg.DataDir, common.DIR_PERM)
	// options need be adjust due to using hdd or sdd, please reference
	// https://github.com/facebook/rocksdb/wiki/RocksDB-Tuning-Guide
	bbto := gorocksdb.NewDefaultBlockBasedTableOptions()
	// use large block to reduce index block size for hdd
	// if using ssd, should use the default value
	bbto.SetBlockSize(1024 * 16)
	// should about 20% less than host RAM
	// http://smalldatum.blogspot.com/2016/09/tuning-rocksdb-block-cache.html
	bbto.SetBlockCache(gorocksdb.NewLRUCache(1024 * 1024 * 1024))
	// for hdd , we nee cache index and filter blocks
	bbto.SetCacheIndexAndFilterBlocks(true)
	filter := gorocksdb.NewBloomFilter(10)
	bbto.SetFilterPolicy(filter)
	opts := gorocksdb.NewDefaultOptions()
	opts.SetBlockBasedTableFactory(bbto)
	opts.SetCreateIfMissing(true)
	opts.SetMaxOpenFiles(-1)
	// keep level0_file_num_compaction_trigger * write_buffer_size = max_bytes_for_level_base to minimize write amplification
	opts.SetWriteBufferSize(1024 * 1024 * 128)
	opts.SetMaxWriteBufferNumber(8)
	opts.SetLevel0FileNumCompactionTrigger(4)
	opts.SetMaxBytesForLevelBase(1024 * 1024 * 1024 * 2)
	opts.SetMinWriteBufferNumberToMerge(2)
	opts.SetTargetFileSizeBase(1024 * 1024 * 128)
	opts.SetMaxBackgroundFlushes(2)
	opts.SetMaxBackgroundCompactions(4)
	opts.SetMinLevelToCompress(3)
	// we use table, so we use prefix seek feature
	opts.SetPrefixExtractor(gorocksdb.NewFixedPrefixTransform(3))
	opts.SetMemtablePrefixBloomSizeRatio(0.1)
	opts.EnableStatistics()
	opts.SetMaxLogFileSize(1024 * 1024 * 32)
	opts.SetLogFileTimeToRoll(3600 * 24 * 3)
	// https://github.com/facebook/mysql-5.6/wiki/my.cnf-tuning
	// rate limiter need to reduce the compaction io

	db := &RockDB{
		cfg:              cfg,
		dbOpts:           opts,
		defaultReadOpts:  cfg.DefaultReadOpts,
		defaultWriteOpts: cfg.DefaultWriteOpts,
		wb:               gorocksdb.NewWriteBatch(),
		backupC:          make(chan *BackupInfo),
		quit:             make(chan struct{}),
	}
	eng, err := gorocksdb.OpenDb(opts, db.GetDataDir())
	if err != nil {
		return nil, err
	}
	db.eng = eng
	os.MkdirAll(db.GetBackupDir(), common.DIR_PERM)

	db.wg.Add(1)
	go func() {
		defer db.wg.Done()
		db.backupLoop()
	}()
	return db, nil
}

func GetBackupDir(base string) string {
	return path.Join(base, "rocksdb_backup")
}

func (r *RockDB) GetBackupBase() string {
	return r.cfg.DataDir
}

func (r *RockDB) GetBackupDir() string {
	return GetBackupDir(r.cfg.DataDir)
}

func (r *RockDB) GetDataDir() string {
	return path.Join(r.cfg.DataDir, "rocksdb")
}

func (r *RockDB) reOpen() error {
	var err error
	r.eng, err = gorocksdb.OpenDb(r.dbOpts, r.GetDataDir())
	return err
}

func (r *RockDB) CompactRange() {
	var rg gorocksdb.Range
	r.eng.CompactRange(rg)
}

func (r *RockDB) Close() {
	close(r.quit)
	r.wg.Wait()
	if r.defaultReadOpts != nil {
		r.defaultReadOpts.Destroy()
	}
	if r.defaultWriteOpts != nil {
		r.defaultWriteOpts.Destroy()
	}
	if r.eng != nil {
		r.eng.Close()
	}
}

func (r *RockDB) SetPerfLevel(level int) {
	// TODO:
}

func (r *RockDB) GetStatistics() string {
	return r.dbOpts.GetStatistics()
}

func (r *RockDB) GetInternalStatus() map[string]interface{} {
	status := make(map[string]interface{})
	bbt := r.dbOpts.GetBlockBasedTableFactory()
	if bbt != nil {
		bc := bbt.GetBlockCache()
		if bc != nil {
			status["block-cache-usage"] = bc.GetUsage()
			status["block-cache-pinned-usage"] = bc.GetPinnedUsage()
		}
	}

	memStr := r.eng.GetProperty("rocksdb.estimate-table-readers-mem")
	status["estimate-table-readers-mem"] = memStr
	memStr = r.eng.GetProperty("rocksdb.cur-size-all-mem-tables")
	status["cur-size-all-mem-tables"] = memStr
	memStr = r.eng.GetProperty("rocksdb.cur-size-active-mem-table")
	status["cur-size-active-mem-tables"] = memStr
	return status
}

func (r *RockDB) GetInternalPropertyStatus(p string) string {
	return r.eng.GetProperty(p)
}

func (r *RockDB) ReadRange(sKey, eKey []byte, maxNum int) chan common.KVRecord {
	retChan := make(chan common.KVRecord, 32)
	go func() {
		it := NewDBRangeLimitIterator(r.eng, sKey, eKey, RangeClose, 0, maxNum, false)
		defer it.Close()
		for it = it; it.Valid(); it.Next() {
			key := it.Key()
			value := it.Value()
			retChan <- common.KVRecord{Key: key, Value: value}
		}
		close(retChan)
	}()
	return retChan
}

type BackupInfo struct {
	backupDir string
	started   chan struct{}
	done      chan struct{}
	rsp       []byte
	err       error
}

func newBackupInfo(dir string) *BackupInfo {
	return &BackupInfo{
		backupDir: dir,
		started:   make(chan struct{}),
		done:      make(chan struct{}),
	}
}

func (self *BackupInfo) WaitReady() {
	select {
	case <-self.started:
	case <-self.done:
	}
}

func (self *BackupInfo) GetResult() ([]byte, error) {
	select {
	case <-self.done:
	}
	return self.rsp, self.err
}

func (r *RockDB) backupLoop() {
	for {
		select {
		case rsp, ok := <-r.backupC:
			if !ok {
				return
			}

			func() {
				defer close(rsp.done)
				log.Printf("begin backup to:%v \n", rsp.backupDir)
				start := time.Now()
				ck, err := gorocksdb.NewCheckpoint(r.eng)
				if err != nil {
					log.Printf("init checkpoint failed: %v", err)
					rsp.err = err
					return
				}
				_, err = os.Stat(rsp.backupDir)
				if !os.IsNotExist(err) {
					log.Printf("checkpoint exist: %v, remove it", rsp.backupDir)
					os.RemoveAll(rsp.backupDir)
				}
				time.AfterFunc(time.Second*2, func() {
					close(rsp.started)
				})
				err = ck.Save(rsp.backupDir)
				if err != nil {
					log.Printf("save checkpoint failed: %v", err)
					rsp.err = err
					return
				}
				cost := time.Now().Sub(start)
				log.Printf("backup done (cost %v), check point to: %v\n", cost.String(), rsp.backupDir)
				// TODO: purge some old checkpoint
				rsp.rsp = []byte(rsp.backupDir)
			}()
		case <-r.quit:
			return
		}
	}
}

func (r *RockDB) Backup(term uint64, index uint64) *BackupInfo {
	fname := GetCheckpointDir(term, index)
	checkpointDir := path.Join(r.GetBackupDir(), fname)
	bi := newBackupInfo(checkpointDir)
	select {
	case r.backupC <- bi:
	default:
		return nil
	}
	return bi
}

func (r *RockDB) IsLocalBackupOK(term uint64, index uint64) (bool, error) {
	backupDir := r.GetBackupDir()
	checkpointDir := GetCheckpointDir(term, index)
	ro := *r.dbOpts
	ro.SetCreateIfMissing(false)
	db, err := gorocksdb.OpenDbForReadOnly(&ro, path.Join(backupDir, checkpointDir), false)
	if err != nil {
		log.Printf("checkpoint open failed: %v", err)
		return false, err
	}
	db.Close()
	return true, nil
}

func copyFile(src, dst string, override bool) error {
	sfi, err := os.Stat(src)
	if err != nil {
		return err
	}
	if !sfi.Mode().IsRegular() {
		return fmt.Errorf("copyfile: non-regular source file %v (%v)", sfi.Name(), sfi.Mode().String())
	}
	_, err = os.Stat(dst)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	} else {
		if !override {
			return nil
		}
	}
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	_, err = io.Copy(out, in)
	if err != nil {
		out.Close()
		return err
	}
	err = out.Sync()
	if err != nil {
		out.Close()
		return err
	}
	return out.Close()
}

func (r *RockDB) Restore(term uint64, index uint64) error {
	// TODO: maybe write meta (snap term and index) and check the meta data in the backup
	backupDir := r.GetBackupDir()
	hasBackup, _ := r.IsLocalBackupOK(term, index)
	if !hasBackup {
		return errors.New("no backup for restore")
	}

	checkpointDir := GetCheckpointDir(term, index)
	start := time.Now()
	log.Printf("begin restore from checkpoint: %v\n", checkpointDir)
	r.eng.Close()
	// 1. remove all files in current db except sst files
	// 2. get the list of sst in checkpoint
	// 3. remove all the sst files not in the checkpoint list
	// 4. copy all files from checkpoint to current db and do not override sst
	matchName := path.Join(r.GetDataDir(), "*")
	nameList, err := filepath.Glob(matchName)
	if err != nil {
		log.Printf("list files failed:  %v\n", err)
		return err
	}
	ckNameList, err := filepath.Glob(path.Join(backupDir, checkpointDir, "*"))
	if err != nil {
		log.Printf("list checkpoint files failed:  %v\n", err)
		return err
	}
	ckSstNameMap := make(map[string]bool)
	for _, fn := range ckNameList {
		if strings.HasSuffix(fn, ".sst") {
			ckSstNameMap[path.Base(fn)] = true
		}
	}

	for _, fn := range nameList {
		shortName := path.Base(fn)
		if strings.HasPrefix(shortName, "LOG") {
			continue
		}
		if strings.HasSuffix(shortName, ".sst") {
			if _, ok := ckSstNameMap[shortName]; ok {
				log.Printf("keeping sst file: %v", fn)
				continue
			}
		}
		log.Printf("removing: %v", fn)
		os.RemoveAll(fn)
	}
	for _, fn := range ckNameList {
		dst := path.Join(r.GetDataDir(), path.Base(fn))
		err := copyFile(fn, dst, false)
		if err != nil {
			log.Printf("copy %v to %v failed: %v", fn, dst, err)
			return err
		} else {
			log.Printf("copy %v to %v done", fn, dst)
		}
	}

	err = r.reOpen()
	log.Printf("restore done, cost: %v\n", time.Now().Sub(start))
	if err != nil {
		log.Printf("reopen the restored db failed:  %v\n", err)
	}
	return err
}

func (r *RockDB) ClearBackup(term uint64, index uint64) error {
	backupDir := r.GetBackupDir()
	checkpointDir := GetCheckpointDir(term, index)
	return os.RemoveAll(path.Join(backupDir, checkpointDir))
}
<|endoftext|>


<commit_before>package blinker

import (
	"fmt"
	"github.com/golang/glog"
	"io"
	"os"
	"os/exec"
	"path/filepath"
)

type serviceImpl struct {
	settings Settings
}

func getPath(root, country, region, id string) string {
	return filepath.Join(root, fmt.Sprintf("%s-%s-%s", country, region, id))
}

func NewService(settings Settings) (Service, error) {

	impl := &serviceImpl{
		settings: settings,
	}
	return impl, nil
}

func (this *serviceImpl) GetImage(country, region, id string) (bytes io.ReadCloser, size int64, err error) {
	path := getPath(this.settings.FsSettings.RootDir, country, region, id)
	glog.Infoln("Reading from file", path)

	f, err := os.Open(path)
	if err != nil {
		return
	}

	stat, err := f.Stat()
	if err != nil {
		return
	}

	bytes = f
	size = stat.Size()
	return
}

func (this *serviceImpl) ExecAlpr(country, region, id string, image io.ReadCloser) (stdout []byte, err error) {
	path := getPath(this.settings.FsSettings.RootDir, country, region, id)

	glog.Infoln("ExecAlpr: saving to file", path)

	dst, err := os.Create(path)
	defer dst.Close()

	if err != nil {
		return
	}

	_, err = io.Copy(dst, image)
	if err != nil {
		return
	}

	alpr := &AlprCommand{
		Country: country,
		Region:  region,
		Path:    path,
	}

	stdout, err = alpr.Execute()
	if err != nil {
		return
	}

	// copy the results
	json, err := os.Create(path + ".json")

	defer json.Close()

	glog.Infoln("ExecAlpr: saving results to", json.Name())
	json.Write(stdout)

	return
}

func (this *serviceImpl) Close() {
	glog.Infoln("Service closed")
}

func (this *AlprCommand) Execute() (stdout []byte, err error) {
	cmd := exec.Command("alpr", "-c", this.Country, "-t", this.Region, "-j", this.Path)
	glog.Infoln("exec command:", cmd)
	stdout, err = cmd.CombinedOutput()
	glog.Infoln("exec result", string(stdout), err)
	return
}
<commit_msg>Adding image transcoding to png if image type is not known<commit_after>package blinker

import (
	"fmt"
	"github.com/golang/glog"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	"image/png"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

type serviceImpl struct {
	settings Settings
}

func getPath(root, country, region, id string) string {
	return filepath.Join(root, fmt.Sprintf("%s-%s-%s", country, region, id))
}

func NewService(settings Settings) (Service, error) {

	impl := &serviceImpl{
		settings: settings,
	}
	return impl, nil
}

func (this *serviceImpl) GetImage(country, region, id string) (bytes io.ReadCloser, size int64, err error) {
	path := getPath(this.settings.FsSettings.RootDir, country, region, id)
	glog.Infoln("Reading from file", path)

	f, err := os.Open(path)
	if err != nil {
		return
	}

	stat, err := f.Stat()
	if err != nil {
		return
	}

	bytes = f
	size = stat.Size()
	return
}

func (this *serviceImpl) ExecAlpr(country, region, id string, src io.ReadCloser) (stdout []byte, err error) {
	path := getPath(this.settings.FsSettings.RootDir, country, region, id)

	glog.Infoln("ExecAlpr: saving to file", path)

	dst, err := os.Create(path)
	defer dst.Close()

	if err != nil {
		return
	}

	_, err = io.Copy(dst, src)
	if err != nil {
		return
	}

	// If we can't tell by the extension
	jsonBase := path
	if filepath.Ext(path) == "" {
		// Transcode the file to png
		glog.Infoln("Transcode to PNG:", path)
		if imgfile, err := os.Open(path); err == nil {
			if img, format, err := image.Decode(imgfile); err == nil {
				glog.Infoln("Image ", path, "is", format)
				if outfile, err := os.Create(path + ".png"); err == nil {
					defer outfile.Close()
					if err := png.Encode(outfile, img); err == nil {
						path = path + ".png"
					}
				}
			}
		}
	} else {
		jsonBase = filepath.Join(filepath.Dir(path), strings.Split(filepath.Base(path), ".")[0])
	}

	alpr := &AlprCommand{
		Country: country,
		Region:  region,
		Path:    path,
	}

	stdout, err = alpr.Execute()
	if err != nil {
		return
	}

	// copy the results
	json, err := os.Create(jsonBase + ".json")

	defer json.Close()

	glog.Infoln("ExecAlpr: saving results to", json.Name())
	json.Write(stdout)

	return
}

func (this *serviceImpl) Close() {
	glog.Infoln("Service closed")
}

func (this *AlprCommand) Execute() (stdout []byte, err error) {
	cmd := exec.Command("alpr", "-c", this.Country, "-t", this.Region, "-j", this.Path)
	glog.Infoln("exec command:", cmd)
	stdout, err = cmd.Output()
	glog.Infoln("exec result", string(stdout), err)
	return
}
<|endoftext|>