<commit_before>package ovssubnet

import (
	"errors"
	"fmt"
	log "github.com/golang/glog"
	"net"
	"time"

	"github.com/openshift/openshift-sdn/ovssubnet/controller/kube"
	"github.com/openshift/openshift-sdn/ovssubnet/controller/lbr"
	"github.com/openshift/openshift-sdn/pkg/api"
	"github.com/openshift/openshift-sdn/pkg/netutils"
)

type OvsController struct {
	subnetRegistry  api.SubnetRegistry
	localIP         string
	localSubnet     *api.Subnet
	hostName        string
	subnetAllocator *netutils.SubnetAllocator
	sig             chan struct{}
	flowController  FlowController
}

type FlowController interface {
	Setup(localSubnet, globalSubnet string) error
	AddOFRules(minionIP, localSubnet, localIP string) error
	DelOFRules(minionIP, localIP string) error
}

func NewKubeController(sub api.SubnetRegistry, hostname string, selfIP string) (*OvsController, error) {
	kubeController, err := NewController(sub, hostname, selfIP)
	if err == nil {
		kubeController.flowController = kube.NewFlowController()
	}
	return kubeController, err
}

func NewDefaultController(sub api.SubnetRegistry, hostname string, selfIP string) (*OvsController, error) {
	defaultController, err := NewController(sub, hostname, selfIP)
	if err == nil {
		defaultController.flowController = lbr.NewFlowController()
	}
	return defaultController, err
}

func NewController(sub api.SubnetRegistry, hostname string, selfIP string) (*OvsController, error) {
	if selfIP == "" {
		addrs, err := net.LookupIP(hostname)
		if err != nil {
			log.Errorf("Failed to lookup IP Address for %s", hostname)
			return nil, err
		}
		for _, addr := range addrs {
			if addr.String() != "127.0.0.1" {
				selfIP = addr.String()
				break
			}
		}
		if selfIP == "" {
			log.Errorf("Failed to lookup valid IP Address for %s (%v)", hostname, addrs)
			return nil, err
		}
	}
	log.Infof("Self IP: %s.", selfIP)
	return &OvsController{
		subnetRegistry:  sub,
		localIP:         selfIP,
		hostName:        hostname,
		localSubnet:     nil,
		subnetAllocator: nil,
		sig:             make(chan struct{}),
	}, nil
}

func (oc *OvsController) StartMaster(sync bool, containerNetwork string, containerSubnetLength uint) error {
	// wait a minute for etcd to come alive
	status := oc.subnetRegistry.CheckEtcdIsAlive(60)
	if !status {
		log.Errorf("Etcd not running?")
		return errors.New("Etcd not reachable. Sync cluster check failed.")
	}
	// initialize the minion key
	if sync {
		err := oc.subnetRegistry.InitMinions()
		if err != nil {
			log.Infof("Minion path already initialized.")
		}
	}

	// initialize the subnet key?
	err := oc.subnetRegistry.InitSubnets()
	subrange := make([]string, 0)
	if err != nil {
		subnets, err := oc.subnetRegistry.GetSubnets()
		if err != nil {
			log.Errorf("Error in initializing/fetching subnets: %v", err)
			return err
		}
		for _, sub := range *subnets {
			subrange = append(subrange, sub.Sub)
		}
	}

	err = oc.subnetRegistry.WriteNetworkConfig(containerNetwork, containerSubnetLength)
	if err != nil {
		return err
	}

	oc.subnetAllocator, err = netutils.NewSubnetAllocator(containerNetwork, containerSubnetLength, subrange)
	if err != nil {
		return err
	}
	err = oc.ServeExistingMinions()
	if err != nil {
		log.Warningf("Error initializing existing minions: %v", err)
		// no worry, we can still keep watching it.
	}
	go oc.watchMinions()
	return nil
}

func (oc *OvsController) ServeExistingMinions() error {
	minions, err := oc.subnetRegistry.GetMinions()
	if err != nil {
		return err
	}

	for _, minion := range *minions {
		_, err := oc.subnetRegistry.GetSubnet(minion)
		if err == nil {
			// subnet already exists, continue
			continue
		}
		err = oc.AddNode(minion)
		if err != nil {
			return err
		}
	}
	return nil
}

func (oc *OvsController) AddNode(minion string) error {
	sn, err := oc.subnetAllocator.GetNetwork()
	if err != nil {
		log.Errorf("Error creating network for minion %s.", minion)
		return err
	}
	var minionIP string
	ip := net.ParseIP(minion)
	if ip == nil {
		addrs, err := net.LookupIP(minion)
		if err != nil {
			log.Errorf("Failed to lookup IP address for minion %s: %v", minion, err)
			return err
		}
		minionIP = addrs[0].String()
		if minionIP == "" {
			return fmt.Errorf("Failed to obtain IP address from minion label: %s", minion)
		}
	} else {
		minionIP = ip.String()
	}
	sub := &api.Subnet{
		Minion: minionIP,
		Sub:    sn.String(),
	}
	oc.subnetRegistry.CreateSubnet(minion, sub)
	if err != nil {
		log.Errorf("Error writing subnet to etcd for minion %s: %v", minion, sn)
		return err
	}
	return nil
}

func (oc *OvsController) DeleteNode(minion string) error {
	sub, err := oc.subnetRegistry.GetSubnet(minion)
	if err != nil {
		log.Errorf("Error fetching subnet for minion %s for delete operation.", minion)
		return err
	}
	_, ipnet, err := net.ParseCIDR(sub.Sub)
	if err != nil {
		log.Errorf("Error parsing subnet for minion %s for deletion: %s", minion, sub.Sub)
		return err
	}
	oc.subnetAllocator.ReleaseNetwork(ipnet)
	return oc.subnetRegistry.DeleteSubnet(minion)
}

func (oc *OvsController) syncWithMaster() error {
	return oc.subnetRegistry.CreateMinion(oc.hostName, oc.localIP)
}

func (oc *OvsController) StartNode(sync, skipsetup bool) error {
	if sync {
		err := oc.syncWithMaster()
		if err != nil {
			log.Errorf("Failed to register with master: %v", err)
			return err
		}
	}
	err := oc.initSelfSubnet()
	if err != nil {
		log.Errorf("Failed to get subnet for this host: %v", err)
		return err
	}
	// call flow controller's setup
	if err == nil {
		if !skipsetup {
			// Assume we are working with IPv4
			containerNetwork, err := oc.subnetRegistry.GetContainerNetwork()
			if err != nil {
				log.Errorf("Failed to obtain ContainerNetwork: %v", err)
				return err
			}
			err = oc.flowController.Setup(oc.localSubnet.Sub, containerNetwork)
			if err != nil {
				return err
			}
		}
		subnets, err := oc.subnetRegistry.GetSubnets()
		if err != nil {
			log.Errorf("Could not fetch existing subnets: %v", err)
		}
		for _, s := range *subnets {
			oc.flowController.AddOFRules(s.Minion, s.Sub, oc.localIP)
		}
		go oc.watchCluster()
	}
	return err
}

func (oc *OvsController) initSelfSubnet() error {
	// get subnet for self
	for {
		sub, err := oc.subnetRegistry.GetSubnet(oc.hostName)
		if err != nil {
			log.Errorf("Could not find an allocated subnet for minion %s: %s. Waiting...", oc.hostName, err)
			time.Sleep(2 * time.Second)
			continue
		}
		oc.localSubnet = sub
		return nil
	}
}

func (oc *OvsController) watchMinions() {
	// watch latest?
	stop := make(chan bool)
	minevent := make(chan *api.MinionEvent)
	go oc.subnetRegistry.WatchMinions(minevent, stop)
	for {
		select {
		case ev := <-minevent:
			switch ev.Type {
			case api.Added:
				_, err := oc.subnetRegistry.GetSubnet(ev.Minion)
				if err != nil {
					// subnet does not exist already
					oc.AddNode(ev.Minion)
				}
			case api.Deleted:
				oc.DeleteNode(ev.Minion)
			}
		case <-oc.sig:
			log.Error("Signal received. Stopping watching of minions.")
			stop <- true
			return
		}
	}
}

func (oc *OvsController) watchCluster() {
	stop := make(chan bool)
	clusterEvent := make(chan *api.SubnetEvent)
	go oc.subnetRegistry.WatchSubnets(clusterEvent, stop)
	for {
		select {
		case ev := <-clusterEvent:
			switch ev.Type {
			case api.Added:
				// add openflow rules
				oc.flowController.AddOFRules(ev.Sub.Minion, ev.Sub.Sub, oc.localIP)
			case api.Deleted:
				// delete openflow rules meant for the minion
				oc.flowController.DelOFRules(ev.Sub.Minion, oc.localIP)
			}
		case <-oc.sig:
			stop <- true
			return
		}
	}
}

func (oc *OvsController) Stop() {
	close(oc.sig)
	//oc.sig <- struct{}{}
}
<commit_msg>Fix origin issue#2834; submit existing subnets as in-use<commit_after>package ovssubnet

import (
	"errors"
	"fmt"
	log "github.com/golang/glog"
	"net"
	"time"

	"github.com/openshift/openshift-sdn/ovssubnet/controller/kube"
	"github.com/openshift/openshift-sdn/ovssubnet/controller/lbr"
	"github.com/openshift/openshift-sdn/pkg/api"
	"github.com/openshift/openshift-sdn/pkg/netutils"
)

type OvsController struct {
	subnetRegistry  api.SubnetRegistry
	localIP         string
	localSubnet     *api.Subnet
	hostName        string
	subnetAllocator *netutils.SubnetAllocator
	sig             chan struct{}
	flowController  FlowController
}

type FlowController interface {
	Setup(localSubnet, globalSubnet string) error
	AddOFRules(minionIP, localSubnet, localIP string) error
	DelOFRules(minionIP, localIP string) error
}

func NewKubeController(sub api.SubnetRegistry, hostname string, selfIP string) (*OvsController, error) {
	kubeController, err := NewController(sub, hostname, selfIP)
	if err == nil {
		kubeController.flowController = kube.NewFlowController()
	}
	return kubeController, err
}

func NewDefaultController(sub api.SubnetRegistry, hostname string, selfIP string) (*OvsController, error) {
	defaultController, err := NewController(sub, hostname, selfIP)
	if err == nil {
		defaultController.flowController = lbr.NewFlowController()
	}
	return defaultController, err
}

func NewController(sub api.SubnetRegistry, hostname string, selfIP string) (*OvsController, error) {
	if selfIP == "" {
		addrs, err := net.LookupIP(hostname)
		if err != nil {
			log.Errorf("Failed to lookup IP Address for %s", hostname)
			return nil, err
		}
		for _, addr := range addrs {
			if addr.String() != "127.0.0.1" {
				selfIP = addr.String()
				break
			}
		}
		if selfIP == "" {
			log.Errorf("Failed to lookup valid IP Address for %s (%v)", hostname, addrs)
			return nil, err
		}
	}
	log.Infof("Self IP: %s.", selfIP)
	return &OvsController{
		subnetRegistry:  sub,
		localIP:         selfIP,
		hostName:        hostname,
		localSubnet:     nil,
		subnetAllocator: nil,
		sig:             make(chan struct{}),
	}, nil
}

func (oc *OvsController) StartMaster(sync bool, containerNetwork string, containerSubnetLength uint) error {
	// wait a minute for etcd to come alive
	status := oc.subnetRegistry.CheckEtcdIsAlive(60)
	if !status {
		log.Errorf("Etcd not running?")
		return errors.New("Etcd not reachable. Sync cluster check failed.")
	}
	// initialize the minion key
	if sync {
		err := oc.subnetRegistry.InitMinions()
		if err != nil {
			log.Infof("Minion path already initialized.")
		}
	}

	// initialize the subnet key?
	oc.subnetRegistry.InitSubnets()
	subrange := make([]string, 0)
	subnets, err := oc.subnetRegistry.GetSubnets()
	if err != nil {
		log.Errorf("Error in initializing/fetching subnets: %v", err)
		return err
	}
	for _, sub := range *subnets {
		subrange = append(subrange, sub.Sub)
	}

	err = oc.subnetRegistry.WriteNetworkConfig(containerNetwork, containerSubnetLength)
	if err != nil {
		return err
	}

	oc.subnetAllocator, err = netutils.NewSubnetAllocator(containerNetwork, containerSubnetLength, subrange)
	if err != nil {
		return err
	}
	err = oc.ServeExistingMinions()
	if err != nil {
		log.Warningf("Error initializing existing minions: %v", err)
		// no worry, we can still keep watching it.
	}
	go oc.watchMinions()
	return nil
}

func (oc *OvsController) ServeExistingMinions() error {
	minions, err := oc.subnetRegistry.GetMinions()
	if err != nil {
		return err
	}

	for _, minion := range *minions {
		_, err := oc.subnetRegistry.GetSubnet(minion)
		if err == nil {
			// subnet already exists, continue
			continue
		}
		err = oc.AddNode(minion)
		if err != nil {
			return err
		}
	}
	return nil
}

func (oc *OvsController) AddNode(minion string) error {
	sn, err := oc.subnetAllocator.GetNetwork()
	if err != nil {
		log.Errorf("Error creating network for minion %s.", minion)
		return err
	}
	var minionIP string
	ip := net.ParseIP(minion)
	if ip == nil {
		addrs, err := net.LookupIP(minion)
		if err != nil {
			log.Errorf("Failed to lookup IP address for minion %s: %v", minion, err)
			return err
		}
		minionIP = addrs[0].String()
		if minionIP == "" {
			return fmt.Errorf("Failed to obtain IP address from minion label: %s", minion)
		}
	} else {
		minionIP = ip.String()
	}
	sub := &api.Subnet{
		Minion: minionIP,
		Sub:    sn.String(),
	}
	oc.subnetRegistry.CreateSubnet(minion, sub)
	if err != nil {
		log.Errorf("Error writing subnet to etcd for minion %s: %v", minion, sn)
		return err
	}
	return nil
}

func (oc *OvsController) DeleteNode(minion string) error {
	sub, err := oc.subnetRegistry.GetSubnet(minion)
	if err != nil {
		log.Errorf("Error fetching subnet for minion %s for delete operation.", minion)
		return err
	}
	_, ipnet, err := net.ParseCIDR(sub.Sub)
	if err != nil {
		log.Errorf("Error parsing subnet for minion %s for deletion: %s", minion, sub.Sub)
		return err
	}
	oc.subnetAllocator.ReleaseNetwork(ipnet)
	return oc.subnetRegistry.DeleteSubnet(minion)
}

func (oc *OvsController) syncWithMaster() error {
	return oc.subnetRegistry.CreateMinion(oc.hostName, oc.localIP)
}

func (oc *OvsController) StartNode(sync, skipsetup bool) error {
	if sync {
		err := oc.syncWithMaster()
		if err != nil {
			log.Errorf("Failed to register with master: %v", err)
			return err
		}
	}
	err := oc.initSelfSubnet()
	if err != nil {
		log.Errorf("Failed to get subnet for this host: %v", err)
		return err
	}
	// call flow controller's setup
	if err == nil {
		if !skipsetup {
			// Assume we are working with IPv4
			containerNetwork, err := oc.subnetRegistry.GetContainerNetwork()
			if err != nil {
				log.Errorf("Failed to obtain ContainerNetwork: %v", err)
				return err
			}
			err = oc.flowController.Setup(oc.localSubnet.Sub, containerNetwork)
			if err != nil {
				return err
			}
		}
		subnets, err := oc.subnetRegistry.GetSubnets()
		if err != nil {
			log.Errorf("Could not fetch existing subnets: %v", err)
		}
		for _, s := range *subnets {
			oc.flowController.AddOFRules(s.Minion, s.Sub, oc.localIP)
		}
		go oc.watchCluster()
	}
	return err
}

func (oc *OvsController) initSelfSubnet() error {
	// get subnet for self
	for {
		sub, err := oc.subnetRegistry.GetSubnet(oc.hostName)
		if err != nil {
			log.Errorf("Could not find an allocated subnet for minion %s: %s. Waiting...", oc.hostName, err)
			time.Sleep(2 * time.Second)
			continue
		}
		oc.localSubnet = sub
		return nil
	}
}

func (oc *OvsController) watchMinions() {
	// watch latest?
	stop := make(chan bool)
	minevent := make(chan *api.MinionEvent)
	go oc.subnetRegistry.WatchMinions(minevent, stop)
	for {
		select {
		case ev := <-minevent:
			switch ev.Type {
			case api.Added:
				_, err := oc.subnetRegistry.GetSubnet(ev.Minion)
				if err != nil {
					// subnet does not exist already
					oc.AddNode(ev.Minion)
				}
			case api.Deleted:
				oc.DeleteNode(ev.Minion)
			}
		case <-oc.sig:
			log.Error("Signal received. Stopping watching of minions.")
			stop <- true
			return
		}
	}
}

func (oc *OvsController) watchCluster() {
	stop := make(chan bool)
	clusterEvent := make(chan *api.SubnetEvent)
	go oc.subnetRegistry.WatchSubnets(clusterEvent, stop)
	for {
		select {
		case ev := <-clusterEvent:
			switch ev.Type {
			case api.Added:
				// add openflow rules
				oc.flowController.AddOFRules(ev.Sub.Minion, ev.Sub.Sub, oc.localIP)
			case api.Deleted:
				// delete openflow rules meant for the minion
				oc.flowController.DelOFRules(ev.Sub.Minion, oc.localIP)
			}
		case <-oc.sig:
			stop <- true
			return
		}
	}
}

func (oc *OvsController) Stop() {
	close(oc.sig)
	//oc.sig <- struct{}{}
}
<|endoftext|>


<commit_before>/*
A simple reverse proxy and load balancer

supported balancing strategies:
	- Round Robin
	- Random
	- Least Connected

Example of a json config:

{
  "servers": [
    {
      "name": "server1",
      "port": 7777,
      "services": [
        {
          "route_prefix": "/microservice/",
          "instances": [
            {
              "host": "127.0.0.1:1111",
              "weight": 3
            },
            {
              "host": "127.0.0.1:2222"
            },
            {
              "host": "127.0.0.1:3333"
            }
          ],
          "balancing_strategy": "round_robin",
          "max_connections": 300
        }
      ]
    }
  ]
}
*/
package pylon

import (
	"net/http"
	"regexp"
	"errors"
	"fmt"
	"strconv"
	"sync"
	"math/rand"
	"net"
	"time"
	"net/http/httputil"
	"strings"
)

var (
	dialer = &net.Dialer{Timeout: 1 * time.Second}
	NoService error = errors.New("Service has no route")
)

type RouteType int8

const (
	maxInt = int(^uint(0) >> 1)
	blacklisted = true

	Regex RouteType = iota
	Prefix RouteType = iota

	defaultMaxCon = 100000
)

type Pylon struct {
	Services []*MicroService
}

type Route interface {
	Type() RouteType
	Data() interface{}
}

type MicroService struct {
	Route       Route
	Instances   []*Instance
	Strategy    Strategy
	//LastUsedIdx chan int
	LastUsedIdx int
	BlackList   map[int]bool
	ReqCount    chan int
	Mutex       *sync.RWMutex
}

type RegexRoute struct {
	Regex *regexp.Regexp
}

type PrefixRoute struct {
	Prefix string
}

func (r RegexRoute) Type() RouteType {
	return Regex
}

func (r RegexRoute) Data() interface{} {
	return r.Regex
}

func (p PrefixRoute) Type() RouteType {
	return Prefix
}

func (p PrefixRoute) Data() interface{} {
	return p.Prefix
}

func ListenAndServe(p string) error {
	jsonParser := JSONConfigParser{}
	c, err := jsonParser.ParseFromPath(p)
	if err != nil {
		return err
	}

	return ListenAndServeConfig(c)
}

func ListenAndServeConfig(c *Config) error {
	wg := sync.WaitGroup{}
	wg.Add(len(c.Servers))

	for _, s := range c.Servers {
		p, err := NewPylon(&s)
		if err != nil {
			return err
		}
		go func() {
			defer wg.Done()
			serve(p, s.Port)
		}()
	}
	wg.Wait()
	return nil
}

func NewPylon(s *Server) (*Pylon, error) {
	p := &Pylon{}
	for _, ser := range s.Services {
		m := MicroService{}
		if ser.Pattern != "" {
			reg, err := regexp.Compile(ser.Pattern)
			if err != nil {
				return nil, err
			}
			m.Route = RegexRoute{
				Regex: reg,
			}
		} else if ser.Prefix != "" {
			m.Route = PrefixRoute{
				Prefix: ser.Prefix,
			}
		} else {
			return nil, NoService
		}

		maxCon := defaultMaxCon
		if ser.MaxCon > 0 {
			maxCon = ser.MaxCon
		}

		for _, inst := range ser.Instances {
			var weight float32 = 1
			if inst.Weight > 0 {
				weight = inst.Weight
			}
			m.Instances = append(m.Instances, &Instance{
				inst.Host,
				weight,
				make(chan int, maxCon),
			})
		}
		m.Strategy = ser.Strategy
		m.BlackList = make(map[int]bool, len(ser.Instances))
		m.Mutex = &sync.RWMutex{}
		//m.LastUsedIdx = make(chan int, 1)
		//m.LastUsedIdx <- 0
		m.ReqCount = make(chan int, maxCon)

		p.Services = append(p.Services, &m)
	}

	return p, nil
}

func serve(p *Pylon, port int) {
	mux := http.NewServeMux()
	mux.Handle("/", NewPylonHandler(p))
	fmt.Println("Serving on " + strconv.Itoa(port))
	server := &http.Server{
		Addr:           ":" + strconv.Itoa(port),
		Handler:        mux,
		ReadTimeout:    20 * time.Second,
		WriteTimeout:   20 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}
	server.ListenAndServe()
}

func NewPylonHandler(p *Pylon) http.HandlerFunc {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
	}
	return func(w http.ResponseWriter, r *http.Request) {
		route, err := getRoute(r.URL.Path, p)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		m := p.microFromRoute(route)
		if m == nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		inst, err := m.getLoadBalancedInst()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		//m.notifyReq(true)
		m.ReqCount <- 1
		inst.ReqCount <- 1
		fmt.Println("Serving new request, count: " + strconv.Itoa(len(m.ReqCount)))
		proxy := httputil.ReverseProxy{
			Director: func(req *http.Request) {
				req.URL.Scheme = "http"
				req.URL.Host = inst.Host
			},
			Transport: transport,
			FlushInterval: 2 * time.Second,
		}
		proxy.ServeHTTP(w, r)
		//r.Body.Close()
		//m.notifyReq(false)
		<-inst.ReqCount
		<-m.ReqCount
		fmt.Println("Request served, count: " + strconv.Itoa(len(m.ReqCount)))
	}
}

func getRoute(path string, p *Pylon) (string, error) {
	for _, ser := range p.Services {
		switch ser.Route.Type() {
		case Regex:
			reg := ser.Route.Data().(*regexp.Regexp)
			if reg.Match([]byte(path)) {
				return reg.String(), nil
			}
		case Prefix:
			pref := ser.Route.Data().(string)
			if strings.HasPrefix(path, pref) {
				return pref, nil
			}
		default: return "", errors.New("Route has no correct type")

		}
	}
	return "no_route", errors.New("No route available for path " + path)
}

func (p *Pylon) microFromRoute(route string) *MicroService {
	for _, ser := range p.Services {
		switch ser.Route.Type() {
		case Regex:
			if ser.Route.Data().(*regexp.Regexp).String() == route {
				return ser
			}
		case Prefix:
			if ser.Route.Data().(string) == route {
				return ser
			}
		default:
			return nil
		}
	}
	return nil
}

func (m *MicroService) getLoadBalancedInst() (*Instance, error) {
	if m == nil {
		return nil, errors.New("No service")
	}

	for {
		inst, _, err := m.getInstance()
		if err != nil {
			return nil, err
		}

		//fmt.Println("Dialing instance " + inst.Host + "...")
		//_, err = dialer.Dial("tcp", inst.Host)
		//if err != nil {
		//	m.blackList(idx)
		//	fmt.Println("Could not contact instance " + inst.Host)
		//	continue
		//}

		fmt.Println("Instance is " + inst.Host)
		return inst, nil
	}
	return nil, fmt.Errorf("No endpoint available for %+v\n", m)
}

func (m *MicroService) getInstance() (*Instance, int, error) {
	instCount := len(m.Instances)
	if instCount == 0 {
		return nil, -1, errors.New("Microservice has no instances")
	}

	if len(m.BlackList) == instCount {
		return nil, -1, errors.New("All instances are dead")
	}

	var idx int
	switch m.Strategy {
	case RoundRobin:
		idx = m.nextRoundRobinInstIdx()
	case LeastConnected:
		idx = m.getLeastConInstIdx()
	case Random:
		idx = m.getRandomInstIdx()
	default:
		return nil, -1, errors.New("Unexpected strategy " + string(m.Strategy))
	}

	if m.isBlacklisted(idx) {
		return m.getInstance()
	} else {
		return m.Instances[idx], idx, nil
	}
}

func (m *MicroService) nextRoundRobinInstIdx() int {
	//m.Mutex.Lock()
	m.LastUsedIdx++
	if m.LastUsedIdx >= len(m.Instances) {
		m.LastUsedIdx = 0
	}
	//m.Mutex.Unlock()

	return m.LastUsedIdx
}

func (m *MicroService) getLeastConInstIdx() int {
	min := maxInt
	idx := 0

	for i, inst := range m.Instances {
		reqCount := len(inst.ReqCount)
		if reqCount < min {
			min = reqCount
			idx = i
		}
	}

	return idx
}

func (m *MicroService) getRandomInstIdx() int {
	var sum float32 = 0.0
	for _, inst := range m.Instances {
		sum += inst.Weight
	}

	r := rand.Float32() * sum

	for i, inst := range m.Instances {
		r -= inst.Weight
		if r < 0 {
			return i
		}
	}

	return 0
}

/*func (m *MicroService) nextInst() int {
	cur := <- m.LastUsedIdx
	cur++
	if cur >= len(m.Instances) {
		cur = 0
	}
	m.LastUsedIdx <- cur
	return cur
}*/

/*func (m *MicroService) notifyReq(in bool) {
	if in {
		m.ReqCount++
	} else {
		m.ReqCount--
	}
}*/

func (m *MicroService) notifyReq(in bool) {
	if in {
		m.ReqCount <- 1
	} else {
		<-m.ReqCount
	}
}

func (m *MicroService) blackList(idx int) {
	m.Mutex.Lock()
	m.BlackList[idx] = blacklisted
	m.Mutex.Unlock()
}

func (m *MicroService) isBlacklisted(idx int) bool {
	blackListed := false
	m.Mutex.RLock()
	blackListed = m.BlackList[idx]
	m.Mutex.RUnlock()

	return blackListed
}

<commit_msg>Change getLeastConInstIdx to take weights into account<commit_after>/*
A simple reverse proxy and load balancer

supported balancing strategies:
	- Round Robin
	- Random
	- Least Connected

Example of a json config:

{
  "servers": [
    {
      "name": "server1",
      "port": 7777,
      "services": [
        {
          "route_prefix": "/microservice/",
          "instances": [
            {
              "host": "127.0.0.1:1111",
              "weight": 3
            },
            {
              "host": "127.0.0.1:2222"
            },
            {
              "host": "127.0.0.1:3333"
            }
          ],
          "balancing_strategy": "round_robin",
          "max_connections": 300
        }
      ]
    }
  ]
}
*/
package pylon

import (
	"net/http"
	"regexp"
	"errors"
	"fmt"
	"strconv"
	"sync"
	"math/rand"
	"net"
	"time"
	"net/http/httputil"
	"strings"
)

var (
	dialer = &net.Dialer{Timeout: 1 * time.Second}
	NoService error = errors.New("Service has no route")
)

type RouteType int8

const (
	maxFloat32 = float32(^uint(0))
	blacklisted = true

	Regex RouteType = iota
	Prefix RouteType = iota

	defaultMaxCon = 100000
)

type Pylon struct {
	Services []*MicroService
}

type Route interface {
	Type() RouteType
	Data() interface{}
}

type MicroService struct {
	Route       Route
	Instances   []*Instance
	Strategy    Strategy
	//LastUsedIdx chan int
	LastUsedIdx int
	BlackList   map[int]bool
	ReqCount    chan int
	Mutex       *sync.RWMutex
}

type RegexRoute struct {
	Regex *regexp.Regexp
}

type PrefixRoute struct {
	Prefix string
}

func (r RegexRoute) Type() RouteType {
	return Regex
}

func (r RegexRoute) Data() interface{} {
	return r.Regex
}

func (p PrefixRoute) Type() RouteType {
	return Prefix
}

func (p PrefixRoute) Data() interface{} {
	return p.Prefix
}

func ListenAndServe(p string) error {
	jsonParser := JSONConfigParser{}
	c, err := jsonParser.ParseFromPath(p)
	if err != nil {
		return err
	}

	return ListenAndServeConfig(c)
}

func ListenAndServeConfig(c *Config) error {
	wg := sync.WaitGroup{}
	wg.Add(len(c.Servers))

	for _, s := range c.Servers {
		p, err := NewPylon(&s)
		if err != nil {
			return err
		}
		go func() {
			defer wg.Done()
			serve(p, s.Port)
		}()
	}
	wg.Wait()
	return nil
}

func NewPylon(s *Server) (*Pylon, error) {
	p := &Pylon{}
	for _, ser := range s.Services {
		m := MicroService{}
		if ser.Pattern != "" {
			reg, err := regexp.Compile(ser.Pattern)
			if err != nil {
				return nil, err
			}
			m.Route = RegexRoute{
				Regex: reg,
			}
		} else if ser.Prefix != "" {
			m.Route = PrefixRoute{
				Prefix: ser.Prefix,
			}
		} else {
			return nil, NoService
		}

		maxCon := defaultMaxCon
		if ser.MaxCon > 0 {
			maxCon = ser.MaxCon
		}

		for _, inst := range ser.Instances {
			var weight float32 = 1
			if inst.Weight > 0 {
				weight = inst.Weight
			}
			m.Instances = append(m.Instances, &Instance{
				inst.Host,
				weight,
				make(chan int, maxCon),
			})
		}
		m.Strategy = ser.Strategy
		m.BlackList = make(map[int]bool, len(ser.Instances))
		m.Mutex = &sync.RWMutex{}
		//m.LastUsedIdx = make(chan int, 1)
		//m.LastUsedIdx <- 0
		m.ReqCount = make(chan int, maxCon)

		p.Services = append(p.Services, &m)
	}

	return p, nil
}

func serve(p *Pylon, port int) {
	mux := http.NewServeMux()
	mux.Handle("/", NewPylonHandler(p))
	fmt.Println("Serving on " + strconv.Itoa(port))
	server := &http.Server{
		Addr:           ":" + strconv.Itoa(port),
		Handler:        mux,
		ReadTimeout:    20 * time.Second,
		WriteTimeout:   20 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}
	server.ListenAndServe()
}

func NewPylonHandler(p *Pylon) http.HandlerFunc {
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
	}
	return func(w http.ResponseWriter, r *http.Request) {
		route, err := getRoute(r.URL.Path, p)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		m := p.microFromRoute(route)
		if m == nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		inst, err := m.getLoadBalancedInst()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		//m.notifyReq(true)
		m.ReqCount <- 1
		inst.ReqCount <- 1
		fmt.Println("Serving new request, count: " + strconv.Itoa(len(m.ReqCount)))
		proxy := httputil.ReverseProxy{
			Director: func(req *http.Request) {
				req.URL.Scheme = "http"
				req.URL.Host = inst.Host
			},
			Transport: transport,
			FlushInterval: 2 * time.Second,
		}
		proxy.ServeHTTP(w, r)
		//r.Body.Close()
		//m.notifyReq(false)
		<-inst.ReqCount
		<-m.ReqCount
		fmt.Println("Request served, count: " + strconv.Itoa(len(m.ReqCount)))
	}
}

func getRoute(path string, p *Pylon) (string, error) {
	for _, ser := range p.Services {
		switch ser.Route.Type() {
		case Regex:
			reg := ser.Route.Data().(*regexp.Regexp)
			if reg.Match([]byte(path)) {
				return reg.String(), nil
			}
		case Prefix:
			pref := ser.Route.Data().(string)
			if strings.HasPrefix(path, pref) {
				return pref, nil
			}
		default: return "", errors.New("Route has no correct type")

		}
	}
	return "no_route", errors.New("No route available for path " + path)
}

func (p *Pylon) microFromRoute(route string) *MicroService {
	for _, ser := range p.Services {
		switch ser.Route.Type() {
		case Regex:
			if ser.Route.Data().(*regexp.Regexp).String() == route {
				return ser
			}
		case Prefix:
			if ser.Route.Data().(string) == route {
				return ser
			}
		default:
			return nil
		}
	}
	return nil
}

func (m *MicroService) getLoadBalancedInst() (*Instance, error) {
	if m == nil {
		return nil, errors.New("No service")
	}

	for {
		inst, _, err := m.getInstance()
		if err != nil {
			return nil, err
		}

		//fmt.Println("Dialing instance " + inst.Host + "...")
		//_, err = dialer.Dial("tcp", inst.Host)
		//if err != nil {
		//	m.blackList(idx)
		//	fmt.Println("Could not contact instance " + inst.Host)
		//	continue
		//}

		fmt.Println("Instance is " + inst.Host)
		return inst, nil
	}
	return nil, fmt.Errorf("No endpoint available for %+v\n", m)
}

func (m *MicroService) getInstance() (*Instance, int, error) {
	instCount := len(m.Instances)
	if instCount == 0 {
		return nil, -1, errors.New("Microservice has no instances")
	}

	if len(m.BlackList) == instCount {
		return nil, -1, errors.New("All instances are dead")
	}

	var idx int
	switch m.Strategy {
	case RoundRobin:
		idx = m.nextRoundRobinInstIdx()
	case LeastConnected:
		idx = m.getLeastConInstIdx()
	case Random:
		idx = m.getRandomInstIdx()
	default:
		return nil, -1, errors.New("Unexpected strategy " + string(m.Strategy))
	}

	if m.isBlacklisted(idx) {
		return m.getInstance()
	} else {
		return m.Instances[idx], idx, nil
	}
}

func (m *MicroService) nextRoundRobinInstIdx() int {
	//m.Mutex.Lock()
	m.LastUsedIdx++
	if m.LastUsedIdx >= len(m.Instances) {
		m.LastUsedIdx = 0
	}
	//m.Mutex.Unlock()

	return m.LastUsedIdx
}

func (m *MicroService) getLeastConInstIdx() int {
	minLoad := maxFloat32
	idx := 0

	for i, inst := range m.Instances {
		load := float32(len(inst.ReqCount)) / inst.Weight
		if load < minLoad {
			minLoad = load
			idx = i
		}
	}

	return idx
}

func (m *MicroService) getRandomInstIdx() int {
	var sum float32 = 0.0
	for _, inst := range m.Instances {
		sum += inst.Weight
	}

	r := rand.Float32() * sum

	for i, inst := range m.Instances {
		r -= inst.Weight
		if r < 0 {
			return i
		}
	}

	return 0
}

/*func (m *MicroService) nextInst() int {
	cur := <- m.LastUsedIdx
	cur++
	if cur >= len(m.Instances) {
		cur = 0
	}
	m.LastUsedIdx <- cur
	return cur
}*/

/*func (m *MicroService) notifyReq(in bool) {
	if in {
		m.ReqCount++
	} else {
		m.ReqCount--
	}
}*/

func (m *MicroService) notifyReq(in bool) {
	if in {
		m.ReqCount <- 1
	} else {
		<-m.ReqCount
	}
}

func (m *MicroService) blackList(idx int) {
	m.Mutex.Lock()
	m.BlackList[idx] = blacklisted
	m.Mutex.Unlock()
}

func (m *MicroService) isBlacklisted(idx int) bool {
	blackListed := false
	m.Mutex.RLock()
	blackListed = m.BlackList[idx]
	m.Mutex.RUnlock()

	return blackListed
}

<|endoftext|>


<commit_before>package dukedb

import (
	"fmt"
	"strings"
	"errors"
)

type Filter interface {
	Type() string
}

/**
 * And.
 */

type AndCondition struct {
	Filters []Filter
}

func (a *AndCondition) Type() string {
	return "and"
}

func And(f ...Filter) *AndCondition {
	a := AndCondition{}
	a.Filters = f
	return &a
}

/**
 * Or.
 */

type OrCondition struct {
 	Filters []Filter
}

func (o *OrCondition) Type() string {
	return "or"
}

func Or(f ...Filter) *OrCondition {
	or := OrCondition{
		Filters: f,
	}
	return &or
}

/**
 * NOT.
 */

type NotCondition struct {
	Filter Filter
}

func (n NotCondition) Type() string {
	return "not"
}

func Not(f Filter) *NotCondition {
	return &NotCondition{
		Filter: f,
	}
}

/**
 * Generic field condition.
 */

type FieldCondition struct {
	Typ string
	Field string
	Value interface{}
}

func (f *FieldCondition) Type() string {
	return f.Typ
}

/**
 * Eq.
 */

func Eq(field string, val interface{}) *FieldCondition {
	eq := FieldCondition{
		Field: field,
		Value: val,
		Typ:"eq",
	}
	return &eq
}

/**
 * Neq.
 */

func Neq(field string, val interface{}) *FieldCondition {
	neq := FieldCondition{
		Field: field,
		Value: val,
		Typ:"neq",
	}
	return &neq
}

/**
 * Like.
 */

func Like(field string, val interface{}) *FieldCondition {
	like := FieldCondition{
		Field: field,
		Value: val,
		Typ:"like",
	}
	return &like
}

/**
 * Less than Lt.
 */

func Lt(field string, val interface{}) *FieldCondition {
	lt := FieldCondition{
		Field: field,
		Value: val,
		Typ:"lt",
	}
	return &lt
}

/**
 * Less than eqal Lte.
 */

func Lte(field string, val interface{}) *FieldCondition {
	lte := FieldCondition{
		Field: field,
		Value: val,
		Typ:"lte",
	}
	return &lte
}

/**
 * Greater than gt.
 */

func Gte(field string, val interface{}) *FieldCondition {
	gte := FieldCondition{
		Field: field,
		Value: val,
		Typ:"gte",
	}
	return &gte
}

func conditionToFilterType(cond string) string {
	typ := ""

	switch strings.ToLower(cond) {
	case "==":
		typ = "eq"
	case "=":
		typ = "eq"
	case "!=":
		typ = "neq"
	case "<":
		typ = "lt"
	case "<=":
		typ = "lte"
	case ">":
		typ = "gt"
	case ">=":
		typ = "gte"
	case "like":
		typ = "like"
	case "in":
		typ = "in"
	default:
		panic(fmt.Sprintf("Unknown field contidion: '%v'", cond))
	}

	return typ
}

/**
 * Query.
 */

type Order struct {
	Field string
	Ascending bool
}

func (o Order) String() string {
	s := o.Field + " "
	if o.Ascending {
		s += "asc"
	} else {
		s += "desc"
	}
	return s
}

type Query struct {
	Model string

	Error error

	Backend Backend

	JoinTargetField string
	JoinField string // The DB name of the field that s joined on. (on the model this query is for)
	JoinedField string // The DB name of the field to join on. (in the parent)

	Joins []*RelationQuery

	LimitNum int
	OffsetNum int
	Orders []Order

	FieldSpec []string
	Filters []Filter
}

func Q(model string) *Query {
	q := Query{}
	q.Model = model

	return &q
}

func (q *Query) Limit(l int) *Query {
	q.LimitNum = l
	return q
}

func (q *Query) Offset(o int) *Query {
	q.OffsetNum = o
	return q
}

func (q *Query) Fields(fields ...string) *Query {
	q.FieldSpec = fields
	return q
}

/**
 * Limit the query to specified fields.
 * If fields where already specified, they will be reduced.
 */
func (q *Query) LimitFields(fields ...string) *Query {
	if q.FieldSpec == nil {
		return q.Fields(fields...)
	}

	allowMap := make(map[string]bool)
	for _, field := range fields {
		allowMap[field] = true
	}

	finalFields := make([]string, 0)

	for _, field := range q.FieldSpec {
		if _, ok := allowMap[field]; ok {
			finalFields = append(finalFields, field)
		}
	}

	q.FieldSpec = finalFields

	return q
}

func (q *Query) Order(name string, asc bool) *Query {
	q.Orders = append(q.Orders, Order{Field: name, Ascending: asc})
	return q
}

func (q *Query) Query(f Filter) *Query {
	q.Filters = append(q.Filters, f)
	return q
}

func (q *Query) filterQ(filter Filter) *Query {
	q.Filters = append(q.Filters, filter)
	return q
}

func (q *Query) Filter(field string, val interface{}) *Query {
	return q.filterQ(Eq(field, val))
}

func (q *Query) FilterCond(field string, condition string, val interface{}) *Query {
	typ := conditionToFilterType(condition)

	f := FieldCondition{
		Typ:typ,
		Field: field,
		Value: val,
	}

	return q.filterQ(&f)
}

func (q *Query) AndQ(filter Filter) *Query {
	return q.filterQ(filter)
}

func (q *Query) And(field string, val interface{}) *Query {
	return q.Filter(field, val)
}

func (q *Query) AndCond(field, condition string, val interface{}) *Query {
	return q.FilterCond(field, condition, val)
}

func (q *Query) OrQ(filter Filter) *Query {
	filterLen := 0
	if q.Filters != nil {
		filterLen = len(q.Filters)
	}

	if filterLen == 0 {
		// No filters set, so just filter regularily.
		return q.filterQ(filter)
	} else if filterLen > 1 {
		// More than one filter.
		// Can not do OR with multiple clauses present.
		q.Error = errors.New("invalid_or_multiple_clauses")
		return q
	}

	// One filter is already present.
	// If it is OR, append to the or.
	// Otherwise create a new top level Or.
	if q.Filters[0].Type() == "or" {
		or := q.Filters[0].(*OrCondition)
		or.Filters = append(or.Filters, filter)
	} else {
		// Other filter is not an OR, so just OR the two together.
		q.Filters = []Filter{Or(q.Filters[0], filter)}
	}

	return q
}

func (q *Query) Or(field string, val interface{}) *Query {
	return q.OrQ(Eq(field, val))
}

func (q *Query) OrCond(field string, condition string, val interface{}) *Query {
	typ := conditionToFilterType(condition)

	f := FieldCondition{
		Typ:typ,
		Field: field,
		Value: val,
	}

	return q.OrQ(&f)
}

/**
 * Joins.
 */

func (q *Query) JoinQ(jq *RelationQuery) *Query {
	jq.BaseQuery = q
	q.Joins = append(q.Joins, jq)
	return q
}

func (q *Query) Join(fieldName string) *Query {
	q.Joins = append(q.Joins, RelQ(q, fieldName))
	return q
}

/**
 * RelationQuery.
 */

func (q *Query) Related(name string) *RelationQuery {
	return RelQ(q, name)
}

func (q *Query) RelatedCustom(name, joinKey, foreignKey string) *RelationQuery {
	return RelQCustom(q, name, joinKey, foreignKey)
}

type RelationQuery struct {
	Query

	BaseQuery *Query
	RelationName string

	JoinFieldName string
	ForeignFieldName string
}

func RelQ(q *Query, name string) *RelationQuery {
	relQ := RelationQuery{
		BaseQuery: q,
		RelationName: name,
	}
	relQ.Backend = q.Backend

	return &relQ
}

func RelQCustom(q *Query, name, joinKey, foreignKey string) *RelationQuery {
	relQ := RelationQuery{
		BaseQuery: q,
		JoinFieldName: joinKey,
		ForeignFieldName: foreignKey,
	}
	relQ.Model = name
	relQ.Backend = q.Backend

	return &relQ
}

func (q *RelationQuery) Find() ([]Model, DbError) {
	if q.Backend == nil {
		panic("Callind .Find() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return nil, err
	}
	return newQ.Find()
}

func (q *RelationQuery) First() (Model, DbError) {
	if q.Backend == nil {
		panic("Callind .First() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return nil, err
	}
	return newQ.First()
}

func (q *RelationQuery) Last() (Model, DbError) {
	if q.Backend == nil {
		panic("Callind .Last() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return nil, err
	}
	return newQ.Last()
}

func (q *RelationQuery) Count() (uint64, DbError) {
	if q.Backend == nil {
		panic("Callind .Count() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return 0, err
	}
	return newQ.Count()
}

func (q *RelationQuery) Delete() DbError {
	if q.Backend == nil {
		panic("Callind .Delete() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return err
	}
	return newQ.Delete()
}

/**
 * Backend functions.
 */

func (q *Query) Find() ([]Model, DbError) {
	if q.Backend == nil {
		panic("Calling .Find() on query without backend")
	}

	return q.Backend.Query(q)
}

func (q *Query) First() (Model, DbError) {
	if q.Backend == nil {
		panic("Calling .First() on query without backend")
	}

	return q.Backend.QueryOne(q)
}

func (q *Query) Last() (Model, DbError) {
	if q.Backend == nil {
		panic("Calling .Last() on query without backend")
	}
	return q.Backend.Last(q)
}

func (q *Query) Count() (uint64, DbError) {
	if q.Backend == nil {
		panic("Calling .Count() on query without backend")
	}
	return q.Backend.Count(q)
}

func (q *Query) Delete() DbError {
	if q.Backend == nil {
		panic("Calling .Delete() on query without backend")
	}
	return q.Backend.DeleteMany(q)
}
<commit_msg>work on query: Renamed Order type to OrderSpec, added SetOrders, SetFilters() and Not(), fixed typos<commit_after>package dukedb

import (
	"fmt"
	"strings"
	"errors"
)

type Filter interface {
	Type() string
}

/**
 * And.
 */

type AndCondition struct {
	Filters []Filter
}

func (a *AndCondition) Type() string {
	return "and"
}

func And(f ...Filter) *AndCondition {
	a := AndCondition{}
	a.Filters = f
	return &a
}

/**
 * Or.
 */

type OrCondition struct {
 	Filters []Filter
}

func (o *OrCondition) Type() string {
	return "or"
}

func Or(f ...Filter) *OrCondition {
	or := OrCondition{
		Filters: f,
	}
	return &or
}

/**
 * NOT.
 */

type NotCondition struct {
	Filter Filter
}

func (n NotCondition) Type() string {
	return "not"
}

func Not(f Filter) *NotCondition {
	return &NotCondition{
		Filter: f,
	}
}

/**
 * Generic field condition.
 */

type FieldCondition struct {
	Typ string
	Field string
	Value interface{}
}

func (f *FieldCondition) Type() string {
	return f.Typ
}

/**
 * Eq.
 */

func Eq(field string, val interface{}) *FieldCondition {
	eq := FieldCondition{
		Field: field,
		Value: val,
		Typ:"eq",
	}
	return &eq
}

/**
 * Neq.
 */

func Neq(field string, val interface{}) *FieldCondition {
	neq := FieldCondition{
		Field: field,
		Value: val,
		Typ:"neq",
	}
	return &neq
}

/**
 * Like.
 */

func Like(field string, val interface{}) *FieldCondition {
	like := FieldCondition{
		Field: field,
		Value: val,
		Typ:"like",
	}
	return &like
}

/**
 * Less than Lt.
 */

func Lt(field string, val interface{}) *FieldCondition {
	lt := FieldCondition{
		Field: field,
		Value: val,
		Typ:"lt",
	}
	return &lt
}

/**
 * Less than eqal Lte.
 */

func Lte(field string, val interface{}) *FieldCondition {
	lte := FieldCondition{
		Field: field,
		Value: val,
		Typ:"lte",
	}
	return &lte
}

/**
 * Greater than gt.
 */

func Gte(field string, val interface{}) *FieldCondition {
	gte := FieldCondition{
		Field: field,
		Value: val,
		Typ:"gte",
	}
	return &gte
}

func conditionToFilterType(cond string) string {
	typ := ""

	switch strings.ToLower(cond) {
	case "==":
		typ = "eq"
	case "=":
		typ = "eq"
	case "!=":
		typ = "neq"
	case "<":
		typ = "lt"
	case "<=":
		typ = "lte"
	case ">":
		typ = "gt"
	case ">=":
		typ = "gte"
	case "like":
		typ = "like"
	case "in":
		typ = "in"
	default:
		panic(fmt.Sprintf("Unknown field contidion: '%v'", cond))
	}

	return typ
}

/**
 * Query.
 */

type OrderSpec struct {
	Field string
	Ascending bool
}

func Order(field string, asc bool) OrderSpec {
	return OrderSpec{Field: field, Ascending: asc}
}

func (o OrderSpec) String() string {
	s := o.Field + " "
	if o.Ascending {
		s += "asc"
	} else {
		s += "desc"
	}
	return s
}

type Query struct {
	Model string

	Error error

	Backend Backend

	JoinTargetField string
	JoinField string // The DB name of the field that s joined on. (on the model this query is for)
	JoinedField string // The DB name of the field to join on. (in the parent)

	Joins []*RelationQuery

	LimitNum int
	OffsetNum int
	Orders []OrderSpec

	FieldSpec []string
	Filters []Filter
}

func Q(model string) *Query {
	q := Query{}
	q.Model = model

	return &q
}

func (q *Query) Limit(l int) *Query {
	q.LimitNum = l
	return q
}

func (q *Query) Offset(o int) *Query {
	q.OffsetNum = o
	return q
}

func (q *Query) Fields(fields ...string) *Query {
	q.FieldSpec = fields
	return q
}

/**
 * Limit the query to specified fields.
 * If fields where already specified, they will be reduced.
 */
func (q *Query) LimitFields(fields ...string) *Query {
	if q.FieldSpec == nil {
		return q.Fields(fields...)
	}

	allowMap := make(map[string]bool)
	for _, field := range fields {
		allowMap[field] = true
	}

	finalFields := make([]string, 0)

	for _, field := range q.FieldSpec {
		if _, ok := allowMap[field]; ok {
			finalFields = append(finalFields, field)
		}
	}

	q.FieldSpec = finalFields

	return q
}

func (q *Query) Order(name string, asc bool) *Query {
	q.Orders = append(q.Orders, OrderSpec{Field: name, Ascending: asc})
	return q
}

func (q *Query) SetOrders(orders ...OrderSpec) *Query {
	q.Orders = orders
	return q
}

func (q *Query) FilterQ(f ...Filter) *Query {
	q.Filters = append(q.Filters, f...)
	return q
}

func (q *Query) SetFilters(f ...Filter) *Query {
	q.Filters = f
	return q
}

func (q *Query) Filter(field string, val interface{}) *Query {
	return q.FilterQ(Eq(field, val))
}

func (q *Query) FilterCond(field string, condition string, val interface{}) *Query {
	typ := conditionToFilterType(condition)

	f := FieldCondition{
		Typ:typ,
		Field: field,
		Value: val,
	}

	return q.FilterQ(&f)
}

func (q *Query) AndQ(filters ...Filter) *Query {
	return q.FilterQ(filters...)
}

func (q *Query) And(field string, val interface{}) *Query {
	return q.Filter(field, val)
}

func (q *Query) AndCond(field, condition string, val interface{}) *Query {
	return q.FilterCond(field, condition, val)
}

func (q *Query) OrQ(filters ...Filter) *Query {
	for _, filter := range filters {
		filterLen := 0
		if q.Filters != nil {
			filterLen = len(q.Filters)
		}

		if filterLen == 0 {
			// No filters set, so just filter regularily.
			return q.FilterQ(filter)
		} else if filterLen > 1 {
			// More than one filter.
			// Can not do OR with multiple clauses present.
			q.Error = errors.New("invalid_or_with_multiple_clauses")
			return q
		}

		// One filter is already present.
		// If it is OR, append to the or.
		// Otherwise create a new top level Or.
		if q.Filters[0].Type() == "or" {
			or := q.Filters[0].(*OrCondition)
			or.Filters = append(or.Filters, filter)
		} else {
			// Other filter is not an OR, so just OR the two together.
			q.Filters = []Filter{Or(q.Filters[0], filter)}
		}
	}

	return q
}

func (q *Query) Or(field string, val interface{}) *Query {
	return q.OrQ(Eq(field, val))
}

func (q *Query) OrCond(field string, condition string, val interface{}) *Query {
	typ := conditionToFilterType(condition)

	f := FieldCondition{
		Typ:typ,
		Field: field,
		Value: val,
	}

	return q.OrQ(&f)
}

func (q *Query) Not(field string, val interface{}) *Query {
	q.Filters = append(q.Filters, Neq(field, val))
	return q
}

/**
 * Joins.
 */

func (q *Query) JoinQ(jq *RelationQuery) *Query {
	jq.BaseQuery = q
	q.Joins = append(q.Joins, jq)
	return q
}

func (q *Query) Join(fieldName string) *Query {
	q.Joins = append(q.Joins, RelQ(q, fieldName))
	return q
}

/**
 * RelationQuery.
 */

func (q *Query) Related(name string) *RelationQuery {
	return RelQ(q, name)
}

func (q *Query) RelatedCustom(name, joinKey, foreignKey string) *RelationQuery {
	return RelQCustom(q, name, joinKey, foreignKey)
}

type RelationQuery struct {
	Query

	BaseQuery *Query
	RelationName string

	JoinFieldName string
	ForeignFieldName string
}

func RelQ(q *Query, name string) *RelationQuery {
	relQ := RelationQuery{
		BaseQuery: q,
		RelationName: name,
	}
	relQ.Backend = q.Backend

	return &relQ
}

func RelQCustom(q *Query, name, joinKey, foreignKey string) *RelationQuery {
	relQ := RelationQuery{
		BaseQuery: q,
		JoinFieldName: joinKey,
		ForeignFieldName: foreignKey,
	}
	relQ.Model = name
	relQ.Backend = q.Backend

	return &relQ
}

func (q *RelationQuery) Find() ([]Model, DbError) {
	if q.Backend == nil {
		panic("Callind .Find() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return nil, err
	}
	return newQ.Find()
}

func (q *RelationQuery) First() (Model, DbError) {
	if q.Backend == nil {
		panic("Calling .First() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return nil, err
	}
	return newQ.First()
}

func (q *RelationQuery) Last() (Model, DbError) {
	if q.Backend == nil {
		panic("Calling .Last() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return nil, err
	}
	return newQ.Last()
}

func (q *RelationQuery) Count() (uint64, DbError) {
	if q.Backend == nil {
		panic("Calling .Count() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return 0, err
	}
	return newQ.Count()
}

func (q *RelationQuery) Delete() DbError {
	if q.Backend == nil {
		panic("Calling .Delete() on a query without backend")
	}

	newQ, err := q.Backend.BuildRelationQuery(q)
	if err != nil {
		return err
	}
	return newQ.Delete()
}

/**
 * Backend functions.
 */

func (q *Query) Find() ([]Model, DbError) {
	if q.Backend == nil {
		panic("Calling .Find() on query without backend")
	}

	return q.Backend.Query(q)
}

func (q *Query) First() (Model, DbError) {
	if q.Backend == nil {
		panic("Calling .First() on query without backend")
	}

	return q.Backend.QueryOne(q)
}

func (q *Query) Last() (Model, DbError) {
	if q.Backend == nil {
		panic("Calling .Last() on query without backend")
	}
	return q.Backend.Last(q)
}

func (q *Query) Count() (uint64, DbError) {
	if q.Backend == nil {
		panic("Calling .Count() on query without backend")
	}
	return q.Backend.Count(q)
}

func (q *Query) Delete() DbError {
	if q.Backend == nil {
		panic("Calling .Delete() on query without backend")
	}
	return q.Backend.DeleteMany(q)
}
<|endoftext|>


<commit_before>package adapter

import (
	"errors"
	"reflect"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/gospackler/caddyshack"
	"github.com/gospackler/couchdb"
)

// View Object is placed over here as a query in couch is possible only with a view
// Helps in forming the javascript that can be used for working on stuff.
type ViewObj struct {
	Name     string
	ViewType reflect.Type
}

// FIXME : Assert Kind of the viewObj passed is a pointer.
func NewViewObj(name string, viewObj caddyshack.StoreObject) *ViewObj {
	// Check if View Exists in the DB.
	// Create View Using the tags if thats the case.
	return &ViewObj{
		Name:     name,
		ViewType: reflect.ValueOf(viewObj).Elem().Type(),
	}
}

// Create the View if it does not exist.
func (v *ViewObj) GetCondition() string {
	// Use the tags to form the condition.
	return "Not implemented"
}

// Initial Version
type CouchQuery struct {
	Condition string // Code for the view RawJson
	ViewName  string
	Store     *CouchStore
	desDoc    *couchdb.DesignDoc
	Params    string
}

func NewQuery(line string, viewName string, desDoc string, db *CouchStore) (couchQuery *CouchQuery) {
	// Assuming a design doc is already created.
	line = "\"map\" : \"" + line + "\""
	return createCouchQuery(line, viewName, desDoc, db, "")
}

func NewMRQuery(mrCode string, viewName string, desDoc string, db *CouchStore, params string) (couchQuery *CouchQuery) {
	return createCouchQuery(mrCode, viewName, desDoc, db, params)
}

func createCouchQuery(rawCond string, viewName string, desDoc string, db *CouchStore, params string) (couchQuery *CouchQuery) {
	desDocObj := db.GetDesignDoc(desDoc)

	couchQuery = &CouchQuery{
		desDoc:    desDocObj,
		Condition: rawCond,
		ViewName:  viewName,
		Store:     db,
		Params:    params,
	}

	// Correct the code over here.
	newView := &couchdb.View{Name: viewName}
	newView.RawJson = couchQuery.Condition

	index, status := couchQuery.desDoc.CheckExists(viewName)

	log.Debug("Views -->", couchQuery.desDoc.Views)
	log.Debug("LastView -->", couchQuery.desDoc.LastView)
	if status == false {
		couchQuery.desDoc.AddView(newView)

		err := couchQuery.desDoc.SaveDoc()
		if err != nil {
			panic(err)
		}

	} else {
		log.Debug("Index found at ", index)
		if index < 0 {
			couchQuery.desDoc.LastView = newView
		} else {
			couchQuery.desDoc.Views[index] = newView
		}
	}
	return
}

// Use reflection to create the query from the tags.
func NewObjQuery(obj caddyshack.StoreObject, db *CouchStore) (q *CouchQuery) {

	q = new(CouchQuery)
	prefix := "doc"

	viewName := q.GetViewName(obj)
	view := couchdb.NewView(viewName, prefix, q.getCondition(obj, prefix), q.getEmits(obj))
	//Creates the DesignDoc if it does not exist.
	desDoc := db.GetDesignDoc(db.Res.DesDoc)
	index, status := desDoc.CheckExists(viewName)

	log.Debug("Index, Status", index, status)

	if status == false {
		desDoc.AddView(view)
		log.Debug("Added view", desDoc, view)
		// FIXME: Removing object update for now. Saved by the guard.
		err := desDoc.SaveDoc()
		if err != nil {
			panic(err)
		}
	} else {
		if index < 0 {
			desDoc.LastView = view
		} else {
			desDoc.Views[index] = view
		}
	}

	q.desDoc = desDoc
	q.ViewName = viewName
	q.Store = db

	return
}

func (q *CouchQuery) SetCondition(cond string) {
	q.Condition = cond
}

func (q *CouchQuery) GetCondition() string {
	return q.Condition
}

func (q *CouchQuery) GetViewName(obj caddyshack.StoreObject) string {
	structObj := reflect.ValueOf(obj).Elem()
	typeOfObj := structObj.Type()
	return strings.ToLower(typeOfObj.Name())
}

func (q *CouchQuery) getCondition(obj caddyshack.StoreObject, prefix string) (condStr string) {

	structObj := reflect.ValueOf(obj).Elem()
	typeOfObj := structObj.Type()

	firstCond := true

	if structObj.Kind() == reflect.Struct {
		for index := 0; index < typeOfObj.NumField(); index++ {
			structField := typeOfObj.Field(index)
			fieldCond := structField.Tag.Get("condition")
			if fieldCond != "" {
				if firstCond {
					condStr = condStr + prefix + "." + fieldCond
					firstCond = false
				} else {
					condStr = condStr + " && " + prefix + "." + fieldCond
				}
			}
		}
	} else {
		panic(errors.New("Expected a struct pointer as input."))
	}

	return
}

func (q *CouchQuery) getEmits(obj caddyshack.StoreObject) (emits string) {

	structObj := reflect.ValueOf(obj).Elem()
	typeOfObj := structObj.Type()

	firstCond := true

	if structObj.Kind() == reflect.Struct {
		for index := 0; index < typeOfObj.NumField(); index++ {
			structField := typeOfObj.Field(index)
			jsonStr := structField.Tag.Get("json")
			if jsonStr != "" {
				if firstCond {
					emits = emits + "\\\"" + jsonStr + "\\\""
					firstCond = false
				} else {
					emits = emits + ", \\\"" + jsonStr + "\\\""
				}
			}
		}
	} else {
		panic(errors.New("Expected a struct pointer as input."))
	}

	return
}

func (q *CouchQuery) Execute() (error, []caddyshack.StoreObject) {
	// Currently O(n) w.r.t to views

	data, err := q.desDoc.Db.GetView(q.desDoc.Id, q.ViewName, q.Params)
	if err != nil {
		return errors.New("Error retreiving view : " + err.Error()), nil
	} else {
		// Print for now create store Object later.
		// FIXME Handle unmarshalling over here.
		result, err := q.Store.MarshalStoreObjects(data)
		if err != nil {
			return errors.New("Could not Marshal json" + err.Error()), result
		}
		return nil, result
	}

	// Move this section to the New
	// The intutive version would be creating an object and then adding methods to it.

	// If it exists get the view back.
	// Otherwise Get Retrieve the Data and Marshal the store Object from the json..
}
<commit_msg>Add a new query with filter capability.<commit_after>package adapter

import (
	"errors"
	"reflect"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/gospackler/caddyshack"
	"github.com/gospackler/couchdb"
)

// View Object is placed over here as a query in couch is possible only with a view
// Helps in forming the javascript that can be used for working on stuff.
type ViewObj struct {
	Name     string
	ViewType reflect.Type
}

// FIXME : Assert Kind of the viewObj passed is a pointer.
func NewViewObj(name string, viewObj caddyshack.StoreObject) *ViewObj {
	// Check if View Exists in the DB.
	// Create View Using the tags if thats the case.
	return &ViewObj{
		Name:     name,
		ViewType: reflect.ValueOf(viewObj).Elem().Type(),
	}
}

// Create the View if it does not exist.
func (v *ViewObj) GetCondition() string {
	// Use the tags to form the condition.
	return "Not implemented"
}

// Initial Version
type CouchQuery struct {
	Condition string // Code for the view RawJson
	ViewName  string
	Store     *CouchStore
	desDoc    *couchdb.DesignDoc
	Params    string
}

func NewQuery(line string, viewName string, desDoc string, db *CouchStore) (couchQuery *CouchQuery) {
	// Assuming a design doc is already created.
	line = "\"map\" : \"" + line + "\""
	return createCouchQuery(line, viewName, desDoc, db, "")
}

func NewFilterQuery(line string, viewName string, desDoc string, db *CouchStore, params string) (couchQuery *CouchQuery) {
	// Assuming a design doc is already created.
	line = "\"map\" : \"" + line + "\""
	return createCouchQuery(line, viewName, desDoc, db, params)
}

func NewMRQuery(mrCode string, viewName string, desDoc string, db *CouchStore, params string) (couchQuery *CouchQuery) {
	return createCouchQuery(mrCode, viewName, desDoc, db, params)
}

func createCouchQuery(rawCond string, viewName string, desDoc string, db *CouchStore, params string) (couchQuery *CouchQuery) {
	desDocObj := db.GetDesignDoc(desDoc)

	couchQuery = &CouchQuery{
		desDoc:    desDocObj,
		Condition: rawCond,
		ViewName:  viewName,
		Store:     db,
		Params:    params,
	}

	// Correct the code over here.
	newView := &couchdb.View{Name: viewName}
	newView.RawJson = couchQuery.Condition

	index, status := couchQuery.desDoc.CheckExists(viewName)

	log.Debug("Views -->", couchQuery.desDoc.Views)
	log.Debug("LastView -->", couchQuery.desDoc.LastView)
	if status == false {
		couchQuery.desDoc.AddView(newView)

		err := couchQuery.desDoc.SaveDoc()
		if err != nil {
			panic(err)
		}

	} else {
		log.Debug("Index found at ", index)
		if index < 0 {
			couchQuery.desDoc.LastView = newView
		} else {
			couchQuery.desDoc.Views[index] = newView
		}
	}
	return
}

// Use reflection to create the query from the tags.
func NewObjQuery(obj caddyshack.StoreObject, db *CouchStore) (q *CouchQuery) {

	q = new(CouchQuery)
	prefix := "doc"

	viewName := q.GetViewName(obj)
	view := couchdb.NewView(viewName, prefix, q.getCondition(obj, prefix), q.getEmits(obj))
	//Creates the DesignDoc if it does not exist.
	desDoc := db.GetDesignDoc(db.Res.DesDoc)
	index, status := desDoc.CheckExists(viewName)

	log.Debug("Index, Status", index, status)

	if status == false {
		desDoc.AddView(view)
		log.Debug("Added view", desDoc, view)
		// FIXME: Removing object update for now. Saved by the guard.
		err := desDoc.SaveDoc()
		if err != nil {
			panic(err)
		}
	} else {
		if index < 0 {
			desDoc.LastView = view
		} else {
			desDoc.Views[index] = view
		}
	}

	q.desDoc = desDoc
	q.ViewName = viewName
	q.Store = db

	return
}

func (q *CouchQuery) SetCondition(cond string) {
	q.Condition = cond
}

func (q *CouchQuery) GetCondition() string {
	return q.Condition
}

func (q *CouchQuery) GetViewName(obj caddyshack.StoreObject) string {
	structObj := reflect.ValueOf(obj).Elem()
	typeOfObj := structObj.Type()
	return strings.ToLower(typeOfObj.Name())
}

func (q *CouchQuery) getCondition(obj caddyshack.StoreObject, prefix string) (condStr string) {

	structObj := reflect.ValueOf(obj).Elem()
	typeOfObj := structObj.Type()

	firstCond := true

	if structObj.Kind() == reflect.Struct {
		for index := 0; index < typeOfObj.NumField(); index++ {
			structField := typeOfObj.Field(index)
			fieldCond := structField.Tag.Get("condition")
			if fieldCond != "" {
				if firstCond {
					condStr = condStr + prefix + "." + fieldCond
					firstCond = false
				} else {
					condStr = condStr + " && " + prefix + "." + fieldCond
				}
			}
		}
	} else {
		panic(errors.New("Expected a struct pointer as input."))
	}

	return
}

func (q *CouchQuery) getEmits(obj caddyshack.StoreObject) (emits string) {

	structObj := reflect.ValueOf(obj).Elem()
	typeOfObj := structObj.Type()

	firstCond := true

	if structObj.Kind() == reflect.Struct {
		for index := 0; index < typeOfObj.NumField(); index++ {
			structField := typeOfObj.Field(index)
			jsonStr := structField.Tag.Get("json")
			if jsonStr != "" {
				if firstCond {
					emits = emits + "\\\"" + jsonStr + "\\\""
					firstCond = false
				} else {
					emits = emits + ", \\\"" + jsonStr + "\\\""
				}
			}
		}
	} else {
		panic(errors.New("Expected a struct pointer as input."))
	}

	return
}

func (q *CouchQuery) Execute() (error, []caddyshack.StoreObject) {
	// Currently O(n) w.r.t to views

	data, err := q.desDoc.Db.GetView(q.desDoc.Id, q.ViewName, q.Params)
	if err != nil {
		return errors.New("Error retreiving view : " + err.Error()), nil
	} else {
		// Print for now create store Object later.
		// FIXME Handle unmarshalling over here.
		result, err := q.Store.MarshalStoreObjects(data)
		if err != nil {
			return errors.New("Could not Marshal json" + err.Error()), result
		}
		return nil, result
	}

	// Move this section to the New
	// The intutive version would be creating an object and then adding methods to it.

	// If it exists get the view back.
	// Otherwise Get Retrieve the Data and Marshal the store Object from the json..
}
<|endoftext|>


<commit_before>package db

import (
	"errors"
	"reflect"
	"testing"

	"github.com/davecgh/go-spew/spew"
	"github.com/microservices-demo/user/users"
)

var (
	TestDB       = fake{}
	ErrFakeError = errors.New("Fake error")
	TestAddress  = users.Address{
		Street:  "street",
		Number:  "51b",
		Country: "Netherlands",
		City:    "Amsterdam",
		ID:      "000056",
	}
)

func TestInit(t *testing.T) {
	err := Init()
	if err == nil {
		t.Error("Expected no registered db error")
	}
	Register("test", TestDB)
	database = "test"
	err = Init()
	if err != ErrFakeError {
		t.Error("expected fake db error from init")
	}
	TestAddress.AddLinks()
}

func TestSet(t *testing.T) {
	database = "nodb"
	err := Set()
	if err == nil {
		t.Error("Expecting error for no databade found")
	}
	Register("nodb2", TestDB)
	database = "nodb2"
	err = Set()
	if err != nil {
		t.Error(err)
	}
}

func TestRegister(t *testing.T) {
	l := len(DBTypes)
	Register("test2", TestDB)
	if len(DBTypes) != l+1 {
		t.Errorf("Expecting %v DB types received %v", l+1, len(DBTypes))
	}
	l = len(DBTypes)
	Register("test2", TestDB)
	if len(DBTypes) != l {
		t.Errorf("Expecting %v DB types received %v duplicate names", l, len(DBTypes))
	}
}

func TestCreateUser(t *testing.T) {
	err := CreateUser(&users.User{})
	if err != ErrFakeError {
		t.Error("expected fake db error from create")
	}
}

func TestGetUser(t *testing.T) {
	_, err := GetUser("test")
	if err != ErrFakeError {
		t.Error("expected fake db error from get")
	}
}

func TestGetUserByName(t *testing.T) {
	_, err := GetUserByName("test")
	if err != ErrFakeError {
		t.Error("expected fake db error from get")
	}
}

func TestGetUserAttributes(t *testing.T) {
	u := users.New()
	GetUserAttributes(&u)
	if len(u.Addresses) != 1 {
		t.Error("expected one address added for GetUserAttributes")
	}
	if !reflect.DeepEqual(u.Addresses[0], TestAddress) {
		spew.Dump(u.Addresses[0])
		t.Error("expected matching addresses")
	}
}

type fake struct{}

func (f fake) Init() error {
	return ErrFakeError
}
func (f fake) GetUserByName(name string) (users.User, error) {
	return users.User{}, ErrFakeError
}
func (f fake) GetUser(id string) (users.User, error) {
	return users.User{}, ErrFakeError
}

func (f fake) GetUsers() ([]users.User, error) {
	return make([]users.User, 0), ErrFakeError
}

func (f fake) CreateUser(*users.User) error {
	return ErrFakeError
}

func (f fake) GetUserAttributes(u *users.User) error {
	u.Addresses = append(u.Addresses, TestAddress)
	return nil
}

func (f fake) GetCard(id string) (users.Card, error) {
	return users.Card{}, ErrFakeError
}

func (f fake) GetCards() ([]users.Card, error) {
	return make([]users.Card, 0), ErrFakeError
}

func (f fake) CreateCard(c *users.Card, id string) error {
	return ErrFakeError
}

func (f fake) GetAddress(id string) (users.Address, error) {
	return users.Address{}, ErrFakeError
}

func (f fake) GetAddresses() ([]users.Address, error) {
	return make([]users.Address, 0), ErrFakeError
}

func (f fake) CreateAddress(u *users.Address, id string) error {
	return ErrFakeError
}
<commit_msg>removed spew<commit_after>package db

import (
	"errors"
	"reflect"
	"testing"

	"github.com/microservices-demo/user/users"
)

var (
	TestDB       = fake{}
	ErrFakeError = errors.New("Fake error")
	TestAddress  = users.Address{
		Street:  "street",
		Number:  "51b",
		Country: "Netherlands",
		City:    "Amsterdam",
		ID:      "000056",
	}
)

func TestInit(t *testing.T) {
	err := Init()
	if err == nil {
		t.Error("Expected no registered db error")
	}
	Register("test", TestDB)
	database = "test"
	err = Init()
	if err != ErrFakeError {
		t.Error("expected fake db error from init")
	}
	TestAddress.AddLinks()
}

func TestSet(t *testing.T) {
	database = "nodb"
	err := Set()
	if err == nil {
		t.Error("Expecting error for no databade found")
	}
	Register("nodb2", TestDB)
	database = "nodb2"
	err = Set()
	if err != nil {
		t.Error(err)
	}
}

func TestRegister(t *testing.T) {
	l := len(DBTypes)
	Register("test2", TestDB)
	if len(DBTypes) != l+1 {
		t.Errorf("Expecting %v DB types received %v", l+1, len(DBTypes))
	}
	l = len(DBTypes)
	Register("test2", TestDB)
	if len(DBTypes) != l {
		t.Errorf("Expecting %v DB types received %v duplicate names", l, len(DBTypes))
	}
}

func TestCreateUser(t *testing.T) {
	err := CreateUser(&users.User{})
	if err != ErrFakeError {
		t.Error("expected fake db error from create")
	}
}

func TestGetUser(t *testing.T) {
	_, err := GetUser("test")
	if err != ErrFakeError {
		t.Error("expected fake db error from get")
	}
}

func TestGetUserByName(t *testing.T) {
	_, err := GetUserByName("test")
	if err != ErrFakeError {
		t.Error("expected fake db error from get")
	}
}

func TestGetUserAttributes(t *testing.T) {
	u := users.New()
	GetUserAttributes(&u)
	if len(u.Addresses) != 1 {
		t.Error("expected one address added for GetUserAttributes")
	}
	if !reflect.DeepEqual(u.Addresses[0], TestAddress) {
		t.Error("expected matching addresses")
	}
}

type fake struct{}

func (f fake) Init() error {
	return ErrFakeError
}
func (f fake) GetUserByName(name string) (users.User, error) {
	return users.User{}, ErrFakeError
}
func (f fake) GetUser(id string) (users.User, error) {
	return users.User{}, ErrFakeError
}

func (f fake) GetUsers() ([]users.User, error) {
	return make([]users.User, 0), ErrFakeError
}

func (f fake) CreateUser(*users.User) error {
	return ErrFakeError
}

func (f fake) GetUserAttributes(u *users.User) error {
	u.Addresses = append(u.Addresses, TestAddress)
	return nil
}

func (f fake) GetCard(id string) (users.Card, error) {
	return users.Card{}, ErrFakeError
}

func (f fake) GetCards() ([]users.Card, error) {
	return make([]users.Card, 0), ErrFakeError
}

func (f fake) CreateCard(c *users.Card, id string) error {
	return ErrFakeError
}

func (f fake) GetAddress(id string) (users.Address, error) {
	return users.Address{}, ErrFakeError
}

func (f fake) GetAddresses() ([]users.Address, error) {
	return make([]users.Address, 0), ErrFakeError
}

func (f fake) CreateAddress(u *users.Address, id string) error {
	return ErrFakeError
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"log"
	"time"
)

var (
	device        = flag.String("device", "lo", "")
	snapshotLen   = flag.Int("snapshot_len", 2048, "")
	BPFFilter     = flag.String("bpf_filter", "tcp and port 5432", "")
	queryFilter   = flag.String("query_filter", "", "not case-sensitive")
	queries       = make(map[string]query)
	slowQueryTime = flag.Int64("slow_query_time", 0, "in milliseconds")
)

type query struct {
	query string
	start time.Time
}

func main() {

	flag.Parse()

	handle, err := pcap.OpenLive(*device, int32(*snapshotLen), true, time.Second)

	defer handle.Close()

	if err != nil {

		log.Fatal(err)
	}

	handle.SetBPFFilter(*BPFFilter)

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())

	var (
		ipLayer  *layers.IPv4
		tcpLayer *layers.TCP
		ok       bool
	)

	for packet := range packetSource.Packets() {

		if applicationLayer := packet.ApplicationLayer(); applicationLayer != nil {

			if ipLayer, ok = packet.Layer(layers.LayerTypeIPv4).(*layers.IPv4); !ok {

				continue
			}

			if tcpLayer, ok = packet.Layer(layers.LayerTypeTCP).(*layers.TCP); !ok {

				continue
			}

			playload := applicationLayer.Payload()

			if len(playload) < 5 {

				continue
			}

			length := _len(playload[1:5])

			if length > len(playload)-1 {

				continue
			}

			switch playload[0] {

			case 'Q', 'P':

				from := fmt.Sprintf("%s%d:%s%d\n", ipLayer.SrcIP, tcpLayer.SrcPort, ipLayer.DstIP, tcpLayer.DstPort)

				if *queryFilter == "" || bytes.Contains(bytes.ToLower(playload[5:length]), bytes.ToLower([]byte(*queryFilter))) {

					queries[from] = query{
						query: string(playload[5:length]),
						start: packet.Metadata().Timestamp,
					}
				}

			default:

				from := fmt.Sprintf("%s%d:%s%d\n", ipLayer.DstIP, tcpLayer.DstPort, ipLayer.SrcIP, tcpLayer.SrcPort)

				if query, found := queries[from]; found {

					queryTime := packet.Metadata().Timestamp.Sub(query.start)

					if *slowQueryTime == 0 || queryTime.Nanoseconds()/1000000 > *slowQueryTime {

						fmt.Printf("-[ QUERY %f s]-:\n%s\n\n\n", queryTime.Seconds(), query.query)
					}

					delete(queries, from)
				}
			}
		}
	}
}

func _len(b []byte) int {

	return int(b[0])<<24 | int(b[1])<<16 | int(b[2])<<8 | int(b[3])
}
<commit_msg>slow time<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"log"
	"time"
)

var (
	device        = flag.String("device", "lo", "")
	snapshotLen   = flag.Int("snapshot_len", 2048, "")
	BPFFilter     = flag.String("bpf_filter", "tcp and port 5432", "")
	queryFilter   = flag.String("query_filter", "", "not case-sensitive")
	queries       = make(map[string]query)
	slowQueryTime = flag.Int64("slow_query_time", 0, "in milliseconds")
)

type query struct {
	query string
	start time.Time
}

func main() {

	flag.Parse()

	handle, err := pcap.OpenLive(*device, int32(*snapshotLen), true, time.Second)

	defer handle.Close()

	if err != nil {

		log.Fatal(err)
	}

	handle.SetBPFFilter(*BPFFilter)

	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())

	var (
		ipLayer  *layers.IPv4
		tcpLayer *layers.TCP
		ok       bool
	)

	for packet := range packetSource.Packets() {

		if applicationLayer := packet.ApplicationLayer(); applicationLayer != nil {

			if ipLayer, ok = packet.Layer(layers.LayerTypeIPv4).(*layers.IPv4); !ok {

				continue
			}

			if tcpLayer, ok = packet.Layer(layers.LayerTypeTCP).(*layers.TCP); !ok {

				continue
			}

			playload := applicationLayer.Payload()

			if len(playload) < 5 {

				continue
			}

			length := _len(playload[1:5])

			if length > len(playload)-1 {

				continue
			}

			switch playload[0] {

			case 'Q', 'P':

				from := fmt.Sprintf("%s%d:%s%d\n", ipLayer.SrcIP, tcpLayer.SrcPort, ipLayer.DstIP, tcpLayer.DstPort)

				if *queryFilter == "" || bytes.Contains(bytes.ToLower(playload[5:length]), bytes.ToLower([]byte(*queryFilter))) {

					queries[from] = query{
						query: string(playload[5:length]),
						start: packet.Metadata().Timestamp,
					}
				}

			default:

				from := fmt.Sprintf("%s%d:%s%d\n", ipLayer.DstIP, tcpLayer.DstPort, ipLayer.SrcIP, tcpLayer.SrcPort)

				if query, found := queries[from]; found {

					queryTime := packet.Metadata().Timestamp.Sub(query.start)

					if *slowQueryTime == 0 || queryTime.Nanoseconds() > *slowQueryTime*1000000 {

						fmt.Printf("-[ QUERY %f s]-:\n%s\n\n\n", queryTime.Seconds(), query.query)
					}

					delete(queries, from)
				}
			}
		}
	}
}

func _len(b []byte) int {

	return int(b[0])<<24 | int(b[1])<<16 | int(b[2])<<8 | int(b[3])
}
<|endoftext|>


<commit_before>package device

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/pkg/errors"
	"golang.org/x/sys/unix"

	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/revert"
	storageDrivers "github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/lxd/storage/filesystem"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/osarch"
	"github.com/lxc/lxd/shared/subprocess"
)

// BlockFsDetect detects the type of block device.
func BlockFsDetect(dev string) (string, error) {
	out, err := shared.RunCommand("blkid", "-s", "TYPE", "-o", "value", dev)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(out), nil
}

// IsBlockdev returns boolean indicating whether device is block type.
func IsBlockdev(path string) bool {
	// Get a stat struct from the provided path
	stat := unix.Stat_t{}
	err := unix.Stat(path, &stat)
	if err != nil {
		return false
	}

	// Check if it's a block device
	if stat.Mode&unix.S_IFMT == unix.S_IFBLK {
		return true
	}

	// Not a device
	return false
}

// DiskMount mounts a disk device.
func DiskMount(srcPath string, dstPath string, readonly bool, recursive bool, propagation string, mountOptions []string, fsName string) error {
	var err error

	// Prepare the mount flags
	flags := 0
	if readonly {
		flags |= unix.MS_RDONLY
	}

	// Detect the filesystem
	if fsName == "none" {
		flags |= unix.MS_BIND
	}

	if propagation != "" {
		switch propagation {
		case "private":
			flags |= unix.MS_PRIVATE
		case "shared":
			flags |= unix.MS_SHARED
		case "slave":
			flags |= unix.MS_SLAVE
		case "unbindable":
			flags |= unix.MS_UNBINDABLE
		case "rprivate":
			flags |= unix.MS_PRIVATE | unix.MS_REC
		case "rshared":
			flags |= unix.MS_SHARED | unix.MS_REC
		case "rslave":
			flags |= unix.MS_SLAVE | unix.MS_REC
		case "runbindable":
			flags |= unix.MS_UNBINDABLE | unix.MS_REC
		default:
			return fmt.Errorf("Invalid propagation mode %q", propagation)
		}
	}

	if recursive {
		flags |= unix.MS_REC
	}

	// Mount the filesystem
	err = unix.Mount(srcPath, dstPath, fsName, uintptr(flags), strings.Join(mountOptions, ","))
	if err != nil {
		return fmt.Errorf("Unable to mount %q at %q with filesystem %q: %w", srcPath, dstPath, fsName, err)
	}

	// Remount bind mounts in readonly mode if requested
	if readonly == true && flags&unix.MS_BIND == unix.MS_BIND {
		flags = unix.MS_RDONLY | unix.MS_BIND | unix.MS_REMOUNT
		err = unix.Mount("", dstPath, fsName, uintptr(flags), "")
		if err != nil {
			return fmt.Errorf("Unable to mount %q in readonly mode: %w", dstPath, err)
		}
	}

	flags = unix.MS_REC | unix.MS_SLAVE
	err = unix.Mount("", dstPath, "", uintptr(flags), "")
	if err != nil {
		return fmt.Errorf("Unable to make mount %q private: %w", dstPath, err)
	}

	return nil
}

// DiskMountClear unmounts and removes the mount path used for disk shares.
func DiskMountClear(mntPath string) error {
	if shared.PathExists(mntPath) {
		if filesystem.IsMountPoint(mntPath) {
			err := unix.Unmount(mntPath, unix.MNT_DETACH)
			if err != nil {
				return errors.Wrapf(err, "Failed unmounting %q", mntPath)
			}
		}

		err := os.Remove(mntPath)
		if err != nil {
			return errors.Wrapf(err, "Failed removing %q", mntPath)
		}
	}

	return nil
}

func diskCephRbdMap(clusterName string, userName string, poolName string, volumeName string) (string, error) {
	devPath, err := shared.RunCommand(
		"rbd",
		"--id", userName,
		"--cluster", clusterName,
		"--pool", poolName,
		"map",
		fmt.Sprintf("%s", volumeName))
	if err != nil {
		return "", err
	}

	idx := strings.Index(devPath, "/dev/rbd")
	if idx < 0 {
		return "", fmt.Errorf("Failed to detect mapped device path")
	}

	devPath = devPath[idx:]
	return strings.TrimSpace(devPath), nil
}

func diskCephRbdUnmap(deviceName string) error {
	unmapImageName := fmt.Sprintf("%s", deviceName)
	busyCount := 0
again:
	_, err := shared.RunCommand(
		"rbd",
		"unmap",
		unmapImageName)
	if err != nil {
		runError, ok := err.(shared.RunError)
		if ok {
			exitError, ok := runError.Err.(*exec.ExitError)
			if ok {
				if exitError.ExitCode() == 22 {
					// EINVAL (already unmapped)
					return nil
				}

				if exitError.ExitCode() == 16 {
					// EBUSY (currently in use)
					busyCount++
					if busyCount == 10 {
						return err
					}

					// Wait a second an try again
					time.Sleep(time.Second)
					goto again
				}
			}
		}

		return err
	}
	goto again
}

func cephFsConfig(clusterName string, userName string) ([]string, string, error) {
	// Get the monitor list.
	monitors, err := storageDrivers.CephMonitors(clusterName)
	if err != nil {
		return nil, "", err
	}

	// Get the keyring entry.
	secret, err := storageDrivers.CephKeyring(clusterName, userName)
	if err != nil {
		return nil, "", err
	}

	return monitors, secret, nil
}

// diskCephfsOptions returns the mntSrcPath and fsOptions to use for mounting a cephfs share.
func diskCephfsOptions(clusterName string, userName string, fsName string, fsPath string) (string, string, error) {
	// Get the credentials and host
	monAddresses, secret, err := cephFsConfig(clusterName, userName)
	if err != nil {
		return "", "", err
	}

	fsOptions := fmt.Sprintf("name=%v,secret=%v,mds_namespace=%v", userName, secret, fsName)
	srcpath := ""
	for _, monAddress := range monAddresses {
		// Add the default port to the mon hosts if not already provided
		if strings.Contains(monAddress, ":6789") {
			srcpath += fmt.Sprintf("%s,", monAddress)
		} else {
			srcpath += fmt.Sprintf("%s:6789,", monAddress)
		}
	}
	srcpath = srcpath[:len(srcpath)-1]
	srcpath += fmt.Sprintf(":/%s", fsPath)

	return srcpath, fsOptions, nil
}

// DiskVMVirtiofsdStart starts a new virtiofsd process.
// Returns UnsupportedError error if the host system or instance does not support virtiosfd, returns normal error
// type if process cannot be started for other reasons.
func DiskVMVirtiofsdStart(inst instance.Instance, socketPath string, pidPath string, logPath string, sharePath string) error {
	revert := revert.New()
	defer revert.Fail()

	// Remove old socket if needed.
	os.Remove(socketPath)

	// Locate virtiofsd.
	cmd, err := exec.LookPath("virtiofsd")
	if err != nil {
		if shared.PathExists("/usr/lib/qemu/virtiofsd") {
			cmd = "/usr/lib/qemu/virtiofsd"
		}
	}

	if cmd == "" {
		return ErrMissingVirtiofsd
	}

	// Currently, virtiofs is broken on at least the ARM architecture.
	// We therefore restrict virtiofs to 64BIT_INTEL_X86.
	if inst.Architecture() != osarch.ARCH_64BIT_INTEL_X86 {
		return UnsupportedError{msg: "Architecture unsupported"}
	}

	if shared.IsTrue(inst.ExpandedConfig()["migration.stateful"]) {
		return UnsupportedError{"Stateful migration unsupported"}
	}

	// Start the virtiofsd process in non-daemon mode.
	proc, err := subprocess.NewProcess(cmd, []string{fmt.Sprintf("--socket-path=%s", socketPath), "-o", fmt.Sprintf("source=%s", sharePath)}, logPath, logPath)
	if err != nil {
		return err
	}

	err = proc.Start()
	if err != nil {
		return errors.Wrapf(err, "Failed to start virtiofsd")
	}

	revert.Add(func() { proc.Stop() })

	err = proc.Save(pidPath)
	if err != nil {
		return errors.Wrapf(err, "Failed to save virtiofsd state")
	}

	// Wait for socket file to exist.
	waitDuration := time.Second * time.Duration(10)
	waitUntil := time.Now().Add(waitDuration)
	for {
		if shared.PathExists(socketPath) {
			break
		}

		if time.Now().After(waitUntil) {
			return fmt.Errorf("virtiofsd failed to bind socket after %v", waitDuration)
		}

		time.Sleep(50 * time.Millisecond)
	}

	revert.Success()
	return nil
}

// DiskVMVirtiofsdStop stops an existing virtiofsd process and cleans up.
func DiskVMVirtiofsdStop(socketPath string, pidPath string) error {
	if shared.PathExists(pidPath) {
		proc, err := subprocess.ImportProcess(pidPath)
		if err != nil {
			return err
		}

		err = proc.Stop()
		// The virtiofsd process will terminate automatically once the VM has stopped.
		// We therefore should only return an error if it's still running and fails to stop.
		if err != nil && err != subprocess.ErrNotRunning {
			return err
		}

		// Remove PID file if needed.
		os.Remove(pidPath)
	}

	// Remove socket file if needed.
	os.Remove(socketPath)

	return nil
}
<commit_msg>lxd/device/device/utils/disk: Update diskCephfsOptions to return []string for mount options<commit_after>package device

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/pkg/errors"
	"golang.org/x/sys/unix"

	"github.com/lxc/lxd/lxd/instance"
	"github.com/lxc/lxd/lxd/revert"
	storageDrivers "github.com/lxc/lxd/lxd/storage/drivers"
	"github.com/lxc/lxd/lxd/storage/filesystem"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/osarch"
	"github.com/lxc/lxd/shared/subprocess"
)

// BlockFsDetect detects the type of block device.
func BlockFsDetect(dev string) (string, error) {
	out, err := shared.RunCommand("blkid", "-s", "TYPE", "-o", "value", dev)
	if err != nil {
		return "", err
	}

	return strings.TrimSpace(out), nil
}

// IsBlockdev returns boolean indicating whether device is block type.
func IsBlockdev(path string) bool {
	// Get a stat struct from the provided path
	stat := unix.Stat_t{}
	err := unix.Stat(path, &stat)
	if err != nil {
		return false
	}

	// Check if it's a block device
	if stat.Mode&unix.S_IFMT == unix.S_IFBLK {
		return true
	}

	// Not a device
	return false
}

// DiskMount mounts a disk device.
func DiskMount(srcPath string, dstPath string, readonly bool, recursive bool, propagation string, mountOptions []string, fsName string) error {
	var err error

	// Prepare the mount flags
	flags := 0
	if readonly {
		flags |= unix.MS_RDONLY
	}

	// Detect the filesystem
	if fsName == "none" {
		flags |= unix.MS_BIND
	}

	if propagation != "" {
		switch propagation {
		case "private":
			flags |= unix.MS_PRIVATE
		case "shared":
			flags |= unix.MS_SHARED
		case "slave":
			flags |= unix.MS_SLAVE
		case "unbindable":
			flags |= unix.MS_UNBINDABLE
		case "rprivate":
			flags |= unix.MS_PRIVATE | unix.MS_REC
		case "rshared":
			flags |= unix.MS_SHARED | unix.MS_REC
		case "rslave":
			flags |= unix.MS_SLAVE | unix.MS_REC
		case "runbindable":
			flags |= unix.MS_UNBINDABLE | unix.MS_REC
		default:
			return fmt.Errorf("Invalid propagation mode %q", propagation)
		}
	}

	if recursive {
		flags |= unix.MS_REC
	}

	// Mount the filesystem
	err = unix.Mount(srcPath, dstPath, fsName, uintptr(flags), strings.Join(mountOptions, ","))
	if err != nil {
		return fmt.Errorf("Unable to mount %q at %q with filesystem %q: %w", srcPath, dstPath, fsName, err)
	}

	// Remount bind mounts in readonly mode if requested
	if readonly == true && flags&unix.MS_BIND == unix.MS_BIND {
		flags = unix.MS_RDONLY | unix.MS_BIND | unix.MS_REMOUNT
		err = unix.Mount("", dstPath, fsName, uintptr(flags), "")
		if err != nil {
			return fmt.Errorf("Unable to mount %q in readonly mode: %w", dstPath, err)
		}
	}

	flags = unix.MS_REC | unix.MS_SLAVE
	err = unix.Mount("", dstPath, "", uintptr(flags), "")
	if err != nil {
		return fmt.Errorf("Unable to make mount %q private: %w", dstPath, err)
	}

	return nil
}

// DiskMountClear unmounts and removes the mount path used for disk shares.
func DiskMountClear(mntPath string) error {
	if shared.PathExists(mntPath) {
		if filesystem.IsMountPoint(mntPath) {
			err := unix.Unmount(mntPath, unix.MNT_DETACH)
			if err != nil {
				return errors.Wrapf(err, "Failed unmounting %q", mntPath)
			}
		}

		err := os.Remove(mntPath)
		if err != nil {
			return errors.Wrapf(err, "Failed removing %q", mntPath)
		}
	}

	return nil
}

func diskCephRbdMap(clusterName string, userName string, poolName string, volumeName string) (string, error) {
	devPath, err := shared.RunCommand(
		"rbd",
		"--id", userName,
		"--cluster", clusterName,
		"--pool", poolName,
		"map",
		fmt.Sprintf("%s", volumeName))
	if err != nil {
		return "", err
	}

	idx := strings.Index(devPath, "/dev/rbd")
	if idx < 0 {
		return "", fmt.Errorf("Failed to detect mapped device path")
	}

	devPath = devPath[idx:]
	return strings.TrimSpace(devPath), nil
}

func diskCephRbdUnmap(deviceName string) error {
	unmapImageName := fmt.Sprintf("%s", deviceName)
	busyCount := 0
again:
	_, err := shared.RunCommand(
		"rbd",
		"unmap",
		unmapImageName)
	if err != nil {
		runError, ok := err.(shared.RunError)
		if ok {
			exitError, ok := runError.Err.(*exec.ExitError)
			if ok {
				if exitError.ExitCode() == 22 {
					// EINVAL (already unmapped)
					return nil
				}

				if exitError.ExitCode() == 16 {
					// EBUSY (currently in use)
					busyCount++
					if busyCount == 10 {
						return err
					}

					// Wait a second an try again
					time.Sleep(time.Second)
					goto again
				}
			}
		}

		return err
	}
	goto again
}

func cephFsConfig(clusterName string, userName string) ([]string, string, error) {
	// Get the monitor list.
	monitors, err := storageDrivers.CephMonitors(clusterName)
	if err != nil {
		return nil, "", err
	}

	// Get the keyring entry.
	secret, err := storageDrivers.CephKeyring(clusterName, userName)
	if err != nil {
		return nil, "", err
	}

	return monitors, secret, nil
}

// diskCephfsOptions returns the mntSrcPath and fsOptions to use for mounting a cephfs share.
func diskCephfsOptions(clusterName string, userName string, fsName string, fsPath string) (string, []string, error) {
	// Get the credentials and host
	monAddresses, secret, err := cephFsConfig(clusterName, userName)
	if err != nil {
		return "", nil, err
	}

	fsOptions := []string{
		fmt.Sprintf("name=%v", userName),
		fmt.Sprintf("secret=%v", secret),
		fmt.Sprintf("mds_namespace=%v", fsName),
	}

	srcpath := ""
	for _, monAddress := range monAddresses {
		// Add the default port to the mon hosts if not already provided
		if strings.Contains(monAddress, ":6789") {
			srcpath += fmt.Sprintf("%s,", monAddress)
		} else {
			srcpath += fmt.Sprintf("%s:6789,", monAddress)
		}
	}
	srcpath = srcpath[:len(srcpath)-1]
	srcpath += fmt.Sprintf(":/%s", fsPath)

	return srcpath, fsOptions, nil
}

// DiskVMVirtiofsdStart starts a new virtiofsd process.
// Returns UnsupportedError error if the host system or instance does not support virtiosfd, returns normal error
// type if process cannot be started for other reasons.
func DiskVMVirtiofsdStart(inst instance.Instance, socketPath string, pidPath string, logPath string, sharePath string) error {
	revert := revert.New()
	defer revert.Fail()

	// Remove old socket if needed.
	os.Remove(socketPath)

	// Locate virtiofsd.
	cmd, err := exec.LookPath("virtiofsd")
	if err != nil {
		if shared.PathExists("/usr/lib/qemu/virtiofsd") {
			cmd = "/usr/lib/qemu/virtiofsd"
		}
	}

	if cmd == "" {
		return ErrMissingVirtiofsd
	}

	// Currently, virtiofs is broken on at least the ARM architecture.
	// We therefore restrict virtiofs to 64BIT_INTEL_X86.
	if inst.Architecture() != osarch.ARCH_64BIT_INTEL_X86 {
		return UnsupportedError{msg: "Architecture unsupported"}
	}

	if shared.IsTrue(inst.ExpandedConfig()["migration.stateful"]) {
		return UnsupportedError{"Stateful migration unsupported"}
	}

	// Start the virtiofsd process in non-daemon mode.
	proc, err := subprocess.NewProcess(cmd, []string{fmt.Sprintf("--socket-path=%s", socketPath), "-o", fmt.Sprintf("source=%s", sharePath)}, logPath, logPath)
	if err != nil {
		return err
	}

	err = proc.Start()
	if err != nil {
		return errors.Wrapf(err, "Failed to start virtiofsd")
	}

	revert.Add(func() { proc.Stop() })

	err = proc.Save(pidPath)
	if err != nil {
		return errors.Wrapf(err, "Failed to save virtiofsd state")
	}

	// Wait for socket file to exist.
	waitDuration := time.Second * time.Duration(10)
	waitUntil := time.Now().Add(waitDuration)
	for {
		if shared.PathExists(socketPath) {
			break
		}

		if time.Now().After(waitUntil) {
			return fmt.Errorf("virtiofsd failed to bind socket after %v", waitDuration)
		}

		time.Sleep(50 * time.Millisecond)
	}

	revert.Success()
	return nil
}

// DiskVMVirtiofsdStop stops an existing virtiofsd process and cleans up.
func DiskVMVirtiofsdStop(socketPath string, pidPath string) error {
	if shared.PathExists(pidPath) {
		proc, err := subprocess.ImportProcess(pidPath)
		if err != nil {
			return err
		}

		err = proc.Stop()
		// The virtiofsd process will terminate automatically once the VM has stopped.
		// We therefore should only return an error if it's still running and fails to stop.
		if err != nil && err != subprocess.ErrNotRunning {
			return err
		}

		// Remove PID file if needed.
		os.Remove(pidPath)
	}

	// Remove socket file if needed.
	os.Remove(socketPath)

	return nil
}
<|endoftext|>


<commit_before>/*

Command gherkinfmt is a command-line gherkin formater and pretty printer.

Usage

	gherkinfmt OPTIONS
	  -centersteps[=false]   formating option, to control step alignment
	  -[no]color             explicitly enable/disable colors
	  -in PATH               path to input file, defaults to stdin
	  -out PATH              path to output file, defaults to stdout

Examples

	$ gherkinfmt -in path/to/some.feature

	$ cat path/to/some.feature | gherkinfmt -centersteps

*/
package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"

	"github.com/muhqu/go-gherkin"
	"github.com/muhqu/go-gherkin/formater"
	"github.com/pebbe/util"
)

var err error
var colors bool
var colorsYes bool
var colorsNo bool
var centerSteps bool
var skipSteps bool
var verbose bool
var inputPath string
var inputReader io.Reader
var outputPath string
var outputWriter io.Writer

func initFlags() {
	flag.Usage = func() {
		self := path.Base(os.Args[0])
		fmt.Fprintf(os.Stderr, `Usage: %[1]s OPTIONS

  -centersteps[=false]   formating option, to control step alignment
  -nosteps[=false]       omit steps, just print scenario headlines

  -[no]color             explicitly enable/disable colors
  -in PATH               path to input file, defaults to stdin
  -out PATH              path to output file, defaults to stdout
  -v                     more verbose error messages

Examples:

  $ %[1]s -in path/to/some.feature

  $ cat path/to/some.feature | %[1]s -centersteps

`, self)
	}
	flag.BoolVar(&colorsYes, "color", false, "explicitly enable colors")
	flag.BoolVar(&colorsNo, "nocolor", false, "explicitly disable colors")
	flag.BoolVar(&centerSteps, "centersteps", false, "formating option, to control step alignment")
	flag.BoolVar(&skipSteps, "nosteps", false, "omit steps, just print scenario headlines")
	flag.BoolVar(&verbose, "v", false, "more verbose error messages")
	flag.StringVar(&inputPath, "in", "", "path to input file")
	flag.StringVar(&outputPath, "out", "", "path to output file")
	flag.Parse()
}

func usageErr(err error) {
	fmt.Fprintf(os.Stderr, "Error: %s\n       Use -h for help.\n", err.Error())
}

func usageErrWithVerboseHint(err error) {
	if verbose {
		usageErr(err)
	} else {
		fmt.Fprintf(os.Stderr, "Error: %s\n       Use -h for help or use -v to increase verbosity\n", err.Error())
	}
}

func main() {
	initFlags()

	if inputPath != "" {
		inputReader, err = os.Open(inputPath)
		if err != nil {
			usageErr(err)
			return
		}
	} else {
		st, err := os.Stdin.Stat()
		if err != nil {
			usageErr(err)
			return
		}
		if st.Size() > 0 {
			inputReader = os.Stdin
		}
	}
	if inputReader == nil {
		usageErr(fmt.Errorf("Missing input (stdin OR -in flag)"))
		return
	}

	if outputPath != "" {
		outputWriter, err = os.Create(outputPath)
		if err != nil {
			usageErr(err)
			return
		}
	} else {
		outputWriter = os.Stdout
	}

	if colorsYes {
		colors = true
	} else if colorsNo {
		colors = false
	} else if outputWriter == os.Stdout {
		colors = util.IsTerminal(os.Stdout)
	}

	fmtr := &formater.GherkinPrettyFormater{
		AnsiColors:  colors,
		CenterSteps: centerSteps,
		SkipSteps:   skipSteps,
	}

	bytes, _ := ioutil.ReadAll(inputReader)
	content := string(bytes)
	gp := gherkin.NewGherkinDOMParser(content)
	gp.Init()
	err = gp.Parse()
	if err != nil {
		usageErrWithVerboseHint(fmt.Errorf("Parsing failed. invalid gherkin"))
		if verbose {
			fmt.Fprintln(os.Stderr, err)
		}
		return
	}
	fmtr.Format(gp, outputWriter)
}
<commit_msg>add cmdline options for new formatter options<commit_after>/*

Command gherkinfmt is a command-line gherkin formater and pretty printer.

Usage

	gherkinfmt OPTIONS

	Formating Options
	  -centersteps[=false]       control step alignment
	  -nosteps[=false]           hide steps
	  -nocomments[=false]        hide comments
	  -noaligncomments[=false]   disable auto alignment of subsequent comments
	  -mincommentindent=NUM      use min indent for aligned comments

	IO Options
	  -[no]color                 explicitly enable/disable colors
	  -in PATH                   path to input file, defaults to stdin
	  -out PATH                  path to output file, defaults to stdout

Examples

	$ gherkinfmt -in path/to/some.feature

	$ cat path/to/some.feature | gherkinfmt -centersteps

*/
package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"

	"github.com/muhqu/go-gherkin"
	"github.com/muhqu/go-gherkin/formater"
	"github.com/pebbe/util"
)

var err error
var colors bool
var colorsYes bool
var colorsNo bool
var centerSteps bool
var skipSteps bool
var skipComments bool
var noCommentAlign bool
var commentAlignMinIndent int
var verbose bool
var inputPath string
var inputReader io.Reader
var outputPath string
var outputWriter io.Writer

func initFlags() {
	flag.Usage = func() {
		self := path.Base(os.Args[0])
		fmt.Fprintf(os.Stderr, `Usage: %[1]s OPTIONS

Formating Options
  -centersteps[=false]       control step alignment
  -nosteps[=false]           hide steps
  -nocomments[=false]        hide comments
  -noaligncomments[=false]   disable auto alignment of subsequent comments
  -mincommentindent=NUM      use min indent for aligned comments

IO Options
  -[no]color                 explicitly enable/disable colors
  -in PATH                   path to input file, defaults to stdin
  -out PATH                  path to output file, defaults to stdout
  -v                         more verbose error messages

Examples:

  $ %[1]s -in path/to/some.feature

  $ cat path/to/some.feature | %[1]s -centersteps

`, self)
	}
	flag.BoolVar(&colorsYes, "color", false, "explicitly enable colors")
	flag.BoolVar(&colorsNo, "nocolor", false, "explicitly disable colors")
	flag.BoolVar(&centerSteps, "centersteps", false, "formating option, to control step alignment")
	flag.BoolVar(&skipSteps, "nosteps", false, "omit steps, just print scenario headlines")
	flag.BoolVar(&skipComments, "nocomments", false, "hide comments")
	flag.BoolVar(&noCommentAlign, "noaligncomments", false, "disable auto alignment of subsequent comments")
	flag.IntVar(&commentAlignMinIndent, "mincommentindent", 45, "use min indent for aligned comments")
	flag.BoolVar(&verbose, "v", false, "more verbose error messages")
	flag.StringVar(&inputPath, "in", "", "path to input file")
	flag.StringVar(&outputPath, "out", "", "path to output file")
	flag.Parse()
}

func usageErr(err error) {
	fmt.Fprintf(os.Stderr, "Error: %s\n       Use -h for help.\n", err.Error())
}

func usageErrWithVerboseHint(err error) {
	if verbose {
		usageErr(err)
	} else {
		fmt.Fprintf(os.Stderr, "Error: %s\n       Use -h for help or use -v to increase verbosity\n", err.Error())
	}
}

func main() {
	initFlags()

	if inputPath != "" {
		inputReader, err = os.Open(inputPath)
		if err != nil {
			usageErr(err)
			return
		}
	} else {
		st, err := os.Stdin.Stat()
		if err != nil {
			usageErr(err)
			return
		}
		if st.Size() > 0 {
			inputReader = os.Stdin
		}
	}
	if inputReader == nil {
		usageErr(fmt.Errorf("Missing input (stdin OR -in flag)"))
		return
	}

	if outputPath != "" {
		outputWriter, err = os.Create(outputPath)
		if err != nil {
			usageErr(err)
			return
		}
	} else {
		outputWriter = os.Stdout
	}

	if colorsYes {
		colors = true
	} else if colorsNo {
		colors = false
	} else if outputWriter == os.Stdout {
		colors = util.IsTerminal(os.Stdout)
	}

	fmtr := &formater.GherkinPrettyFormater{
		AnsiColors:  colors,
		CenterSteps: centerSteps,
		SkipSteps:   skipSteps,
	}

	bytes, _ := ioutil.ReadAll(inputReader)
	content := string(bytes)
	gp := gherkin.NewGherkinDOMParser(content)
	gp.Init()
	err = gp.Parse()
	if err != nil {
		usageErrWithVerboseHint(fmt.Errorf("Parsing failed. invalid gherkin"))
		if verbose {
			fmt.Fprintln(os.Stderr, err)
		}
		return
	}
	fmtr.Format(gp, outputWriter)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/commands/admin"
	"github.com/urfave/cli"
)

func generateHelp([]cli.Command) {
	for idx, cmd := range commands {
		switch cmd.Name {
		case "admin":
			commands[idx].Description = cmd.Description + app.GenerateCommandsHelp(admin.Commands)
		case "commands":
			commands[idx].Description = cmd.Description + app.GenerateCommandsHelp(commands)
		}
	}
}
<commit_msg>Fix admin and commands commands - sneaky bug introduced by !213<commit_after>package main

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/commands/admin"
	"github.com/urfave/cli"
)

func generateHelp(commands []cli.Command) {
	for idx, cmd := range commands {
		switch cmd.Name {
		case "admin":
			commands[idx].Description = cmd.Description + app.GenerateCommandsHelp(admin.Commands)
		case "commands":
			commands[idx].Description = cmd.Description + app.GenerateCommandsHelp(commands)
		}
	}
}
<|endoftext|>