<commit_before>/*
   PulseHA - HA Cluster Daemon
   Copyright (C) 2017-2019  Andrew Zak <andrew@linux.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
package server

import (
	"context"
	log "github.com/Sirupsen/logrus"
	"github.com/Syleron/PulseHA/proto"
	"google.golang.org/grpc/peer"
	"net"
	"runtime"
	"time"
)

/**
Networking - Bring up the groups on the current node
*/
func MakeLocalActive() error {
	log.Debug("Utils:MakeMemberActive() Local node now active")
	for name, node := range DB.Config.Nodes {
		if name == DB.Config.GetLocalNode() {
			for iface, groups := range node.IPGroups {
				for _, groupName := range groups {
					makeGroupActive(iface, groupName)
				}
			}
		}
	}
	return nil
}

/**
Networking - Bring down the ip groups on the current node
*/
func MakeLocalPassive() error {
	log.Debug("Utils:MakeMemberPassive() Local node now passive")
	for name, node := range DB.Config.Nodes {
		if name == DB.Config.GetLocalNode() {
			for iface, groups := range node.IPGroups {
				for _, groupName := range groups {
					makeGroupPassive(iface, groupName)
				}
			}
		}
	}
	return nil
}

/**
Bring up an []ips for a specific interface
*/
func BringUpIPs(iface string, ips []string) error {
	plugin := DB.Plugins.GetNetworkingPlugin()
	if plugin == nil {
		DB.Logging.Debug("Utils:BringUpIps() No networking plugin.. skipping network action")
		return nil
	}
	err := plugin.Plugin.(PluginNet).BringUpIPs(iface, ips)
	return err
}

/**
Bring down an []ips for a specific interface
*/
func BringDownIPs(iface string, ips []string) error {
	plugin := DB.Plugins.GetNetworkingPlugin()
	if plugin == nil {
		DB.Logging.Debug("Utils:BringDownIps() No networking plugin.. skipping network action")
		return nil
	}
	err := plugin.Plugin.(PluginNet).BringDownIPs(iface, ips)
	return err
}

/**

 */
func MyCaller() string {
	// we get the callers as uintptrs - but we just need 1
	fpcs := make([]uintptr, 1)
	// skip 3 levels to get to the caller of whoever called Caller()
	n := runtime.Callers(3, fpcs)
	if n == 0 {
		return "n/a" // proper error her would be better
	}
	// get the info of the actual function that's in the pointer
	fun := runtime.FuncForPC(fpcs[0] - 1)
	if fun == nil {
		return "n/a"
	}
	// return its name
	return fun.Name()
}

/**
Determine who is the correct active node if more than one active is brought online
*/
func GetFailOverCountWinner(members []*proto.MemberlistMember) string {
	for _, member := range members {
		if member.Status != proto.MemberStatus_UNAVAILABLE {
			tym, _ := time.Parse(time.RFC1123, member.LastReceived)
			if tym == (time.Time{}) {
				return member.Hostname
			}
		}
	}
	return ""
}

/**
Determine if a connection is coming in is a member of our config
 */
func CanCommunicate(ctx context.Context) (bool) {
	pr, ok := peer.FromContext(ctx)
	if !ok {
		DB.Logging.Warn("Unable to get address details for context")
		return false
	}
	// check to make sure the peer IP
	_, err := DB.Config.GetNodeHostnameByAddress(pr.Addr.(*net.TCPAddr).IP.String())
	if err != nil {
		DB.Logging.Warn(err.Error() + ". Communication received from another node not in cluster")
		return false
	}
	return true
}
<commit_msg>Syntax updates<commit_after>/*
   PulseHA - HA Cluster Daemon
   Copyright (C) 2017-2019  Andrew Zak <andrew@linux.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
package server

import (
	"context"
	log "github.com/Sirupsen/logrus"
	"github.com/Syleron/PulseHA/proto"
	"google.golang.org/grpc/peer"
	"net"
	"runtime"
	"time"
)

/**
Networking - Bring up the groups on the current node
*/
func MakeLocalActive() error {
	log.Debug("Utils:MakeMemberActive() Local node now active")
	for name, node := range DB.Config.Nodes {
		if name == DB.Config.GetLocalNode() {
			for iface, groups := range node.IPGroups {
				for _, groupName := range groups {
					makeGroupActive(iface, groupName)
				}
			}
		}
	}
	return nil
}

/**
Networking - Bring down the ip groups on the current node
*/
func MakeLocalPassive() error {
	log.Debug("Utils:MakeMemberPassive() Local node now passive")
	for name, node := range DB.Config.Nodes {
		if name == DB.Config.GetLocalNode() {
			for iface, groups := range node.IPGroups {
				for _, groupName := range groups {
					makeGroupPassive(iface, groupName)
				}
			}
		}
	}
	return nil
}

/**
Bring up an []ips for a specific interface
*/
func BringUpIPs(iface string, ips []string) error {
	plugin := DB.Plugins.GetNetworkingPlugin()
	if plugin == nil {
		DB.Logging.Debug("Utils:BringUpIps() No networking plugin.. skipping network action")
		return nil
	}
	err := plugin.Plugin.(PluginNet).BringUpIPs(iface, ips)
	return err
}

/**
Bring down an []ips for a specific interface
*/
func BringDownIPs(iface string, ips []string) error {
	plugin := DB.Plugins.GetNetworkingPlugin()
	if plugin == nil {
		DB.Logging.Debug("Utils:BringDownIps() No networking plugin.. skipping network action")
		return nil
	}
	err := plugin.Plugin.(PluginNet).BringDownIPs(iface, ips)
	return err
}

/**

 */
func MyCaller() string {
	// we get the callers as uintptrs - but we just need 1
	fpcs := make([]uintptr, 1)
	// skip 3 levels to get to the caller of whoever called Caller()
	n := runtime.Callers(3, fpcs)
	if n == 0 {
		return "n/a" // proper error her would be better
	}
	// get the info of the actual function that's in the pointer
	fun := runtime.FuncForPC(fpcs[0] - 1)
	if fun == nil {
		return "n/a"
	}
	// return its name
	return fun.Name()
}

/**
Determine who is the correct active node if more than one active is brought online
*/
func GetFailOverCountWinner(members []*proto.MemberlistMember) string {
	for _, member := range members {
		if member.Status != proto.MemberStatus_UNAVAILABLE {
			tym, _ := time.Parse(time.RFC1123, member.LastReceived)
			if tym == (time.Time{}) {
				return member.Hostname
			}
		}
	}
	return ""
}

/**
Determine if a connection is coming in is a member of our config
 */
func CanCommunicate(ctx context.Context) bool {
	pr, ok := peer.FromContext(ctx)
	if !ok {
		DB.Logging.Warn("Unable to get address details for context")
		return false
	}
	// check to make sure the peer IP
	_, err := DB.Config.GetNodeHostnameByAddress(pr.Addr.(*net.TCPAddr).IP.String())
	if err != nil {
		DB.Logging.Warn(err.Error() + ". Communication received from another node not in cluster")
		return false
	}
	return true
}
<|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,
          "health_check": {
            "enabled": true,
            "interval": 30
          }
        }
      ]
    }
  ]
}
*/
package pylon

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

var (
	dialer = &net.Dialer{Timeout: dialerTimeout}
	NoService error = errors.New("Service has no route")
	proxyPool *ProxyPool
)

type RouteType int8

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

	Regex RouteType = iota
	Prefix RouteType = iota

	defaultMaxCon = 100000
	defaultHealthCheckInterval = 20
	defaultProxyPoolCapacity = defaultMaxCon

	flushInterval = time.Second * 1
	dialerTimeout = time.Second * 3
)

type Pylon struct {
	Services []*MicroService
}

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

type MicroService struct {
	Route       Route
	Instances   []*Instance
	Strategy    Strategy
	LastUsedIdx SharedInt
	BlackList   map[int]bool
	ReqCount    chan int
	// Caching the weight sum for faster retrieval
	WeightSum   float32
	Mutex       *sync.RWMutex
	HealthCheck HealthCheck
}

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{}

	// Initializing the pool before hand in case one server
	// Gets a request as soon as it's served
	poolSize := 0
	for _, s := range c.Servers {
		for _, ser := range s.Services {
			if ser.MaxCon == 0 {
				poolSize += defaultProxyPoolCapacity
			} else {
				poolSize += ser.MaxCon
			}
		}
	}
	fmt.Println("Pool size is", poolSize)
	proxyPool = NewProxyPool(poolSize)

	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, err := NewMicroService(&ser)
		if err != nil {
			return nil, err
		}
		p.Services = append(p.Services, m)
	}

	return p, nil
}

func NewMicroService(s *Service) (*MicroService, error) {
	m := &MicroService{}
	if s.Pattern != "" {
		reg, err := regexp.Compile(s.Pattern)
		if err != nil {
			return nil, err
		}
		m.Route = RegexRoute{
			Regex: reg,
		}
	} else if s.Prefix != "" {
		m.Route = PrefixRoute{
			Prefix: s.Prefix,
		}
	} else {
		return nil, NoService
	}

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

	var weightSum float32 = 0.0
	for _, inst := range s.Instances {
		var weight float32 = 1
		if inst.Weight > 0 {
			weight = inst.Weight
		}
		weightSum += weight

		newInst := &Instance{
			inst.Host,
			weight,
			make(chan int, maxCon),
			NewSharedInt(0),
		}
		m.Instances = append(m.Instances, newInst)
	}
	m.Strategy = s.Strategy
	m.Mutex = &sync.RWMutex{}
	m.BlackList = make(map[int]bool, len(s.Instances))
	m.LastUsedIdx = NewSharedInt(0)
	m.ReqCount = make(chan int, maxCon)
	m.WeightSum = weightSum
	m.HealthCheck = s.HealthCheck
	if m.HealthCheck.Interval == 0 {
		m.HealthCheck.Interval = defaultHealthCheckInterval
	}

	return m, 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,
	}

	for _, s := range p.Services {
		// Do an initial health check
		go handleHealthCheck(s)

		if s.HealthCheck.Enabled {
			go startPeriodicHealthCheck(s, time.Second * time.Duration(s.HealthCheck.Interval))
		}
	}

	server.ListenAndServe()
}

func startPeriodicHealthCheck(m *MicroService, interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for t := range ticker.C {
		fmt.Println("Checking health of Service:", m.Route, " ---tick:", t)
		handleHealthCheck(m)
	}
}

func handleHealthCheck(m *MicroService) bool {
	change := false
	//var weightSum float32 = 0.0
	for i, inst := range m.Instances {
		_, err := dialer.Dial("tcp", inst.Host)
		if err != nil {
			if !m.isBlacklisted(i) {
				m.blackList(i, true)
				change = true
			}
		} else {
			if m.isBlacklisted(i) {
				m.blackList(i, false)
				change = true
			}
			//weightSum += inst.Weight
		}
	}
	//m.Mutex.Lock()
	//m.WeightSum = weightSum
	//m.Mutex.Unlock()
	return change
}

func NewPylonHandler(p *Pylon) http.HandlerFunc {
	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.getLoadBalancedInstance()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		fmt.Println("Instance is " + inst.Host)

		m.ReqCount <- 1
		inst.ReqCount <- 1

		fmt.Println("Serving new request, count: " + strconv.Itoa(len(m.ReqCount)))
		proxy := proxyPool.Get()
		SetUpProxy(proxy, m, inst.Host)
		proxy.ServeHTTP(w, r)
		proxyPool.Put(proxy)

		<-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) getLoadBalancedInstance() (*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("No instance is available")
	}

	instances := make([]*Instance, instCount)
	copy(instances, m.Instances)

	var idx int
	var err error
	for {
		if len(instances) == 0 {
			return nil, -1, errors.New("All instances are dead")
		}

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

		if err != nil {
			return nil, -1, err
		}

		if m.isBlacklisted(idx) {
			instances[idx] = nil
		} else {
			m.LastUsedIdx.Set(idx)
			return instances[idx], idx, nil
		}
	}
}

func nextRoundRobinInstIdx(instances []*Instance, idx int) (int, error) {
	tryCount := 1
	instCount := len(instances)
	lastNonNil := -1
	for {
		inst := instances[idx]
		if inst != nil {
			if inst.isRoundRobinPicked() {
				break
			}
			lastNonNil = idx
		}
		idx++
		tryCount++
		if tryCount > instCount {
			if lastNonNil != -1 {
				return lastNonNil, nil
			}
			return -1, errors.New("No instance can be round robin picked")
		}
		if idx >= instCount {
			idx = 0
		}
	}
	return idx, nil
}

func getLeastConInstIdx(instances []*Instance) int {
	minLoad := maxFloat32
	idx := 0

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

	return idx
}

func getRandomInstIdx(instances []*Instance) int {
	var weightSum float32 = 0.0
	for _, inst := range instances {
		if inst == nil {
			continue
		}
		weightSum += inst.Weight
	}

	r := rand.Float32() * weightSum

	for i, inst := range instances {
		if inst == nil {
			continue
		}
		r -= inst.Weight
		if r < 0 {
			return i
		}
	}

	return 0
}

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

func (m *MicroService) blackList(idx int, blacklist bool) {
	m.Mutex.Lock()
	if blacklist {
		m.BlackList[idx] = blacklisted
	} else {
		delete(m.BlackList, idx)
	}
	m.Mutex.Unlock()
}

func (m *MicroService) blackListHost(host string, blacklist bool) {
	for idx, inst := range m.Instances {
		if inst.Host == host {
			m.blackList(idx, blacklist)
		}
	}
}

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

	return blackListed
}<commit_msg>Fix useless check<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,
          "health_check": {
            "enabled": true,
            "interval": 30
          }
        }
      ]
    }
  ]
}
*/
package pylon

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

var (
	dialer = &net.Dialer{Timeout: dialerTimeout}
	NoService error = errors.New("Service has no route")
	proxyPool *ProxyPool
)

type RouteType int8

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

	Regex RouteType = iota
	Prefix RouteType = iota

	defaultMaxCon = 100000
	defaultHealthCheckInterval = 20
	defaultProxyPoolCapacity = defaultMaxCon

	flushInterval = time.Second * 1
	dialerTimeout = time.Second * 3
)

type Pylon struct {
	Services []*MicroService
}

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

type MicroService struct {
	Route       Route
	Instances   []*Instance
	Strategy    Strategy
	LastUsedIdx SharedInt
	BlackList   map[int]bool
	ReqCount    chan int
	// Caching the weight sum for faster retrieval
	WeightSum   float32
	Mutex       *sync.RWMutex
	HealthCheck HealthCheck
}

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{}

	// Initializing the pool before hand in case one server
	// Gets a request as soon as it's served
	poolSize := 0
	for _, s := range c.Servers {
		for _, ser := range s.Services {
			if ser.MaxCon == 0 {
				poolSize += defaultProxyPoolCapacity
			} else {
				poolSize += ser.MaxCon
			}
		}
	}
	fmt.Println("Pool size is", poolSize)
	proxyPool = NewProxyPool(poolSize)

	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, err := NewMicroService(&ser)
		if err != nil {
			return nil, err
		}
		p.Services = append(p.Services, m)
	}

	return p, nil
}

func NewMicroService(s *Service) (*MicroService, error) {
	m := &MicroService{}
	if s.Pattern != "" {
		reg, err := regexp.Compile(s.Pattern)
		if err != nil {
			return nil, err
		}
		m.Route = RegexRoute{
			Regex: reg,
		}
	} else if s.Prefix != "" {
		m.Route = PrefixRoute{
			Prefix: s.Prefix,
		}
	} else {
		return nil, NoService
	}

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

	var weightSum float32 = 0.0
	for _, inst := range s.Instances {
		var weight float32 = 1
		if inst.Weight > 0 {
			weight = inst.Weight
		}
		weightSum += weight

		newInst := &Instance{
			inst.Host,
			weight,
			make(chan int, maxCon),
			NewSharedInt(0),
		}
		m.Instances = append(m.Instances, newInst)
	}
	m.Strategy = s.Strategy
	m.Mutex = &sync.RWMutex{}
	m.BlackList = make(map[int]bool, len(s.Instances))
	m.LastUsedIdx = NewSharedInt(0)
	m.ReqCount = make(chan int, maxCon)
	m.WeightSum = weightSum
	m.HealthCheck = s.HealthCheck
	if m.HealthCheck.Interval == 0 {
		m.HealthCheck.Interval = defaultHealthCheckInterval
	}

	return m, 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,
	}

	for _, s := range p.Services {
		// Do an initial health check
		go handleHealthCheck(s)

		if s.HealthCheck.Enabled {
			go startPeriodicHealthCheck(s, time.Second * time.Duration(s.HealthCheck.Interval))
		}
	}

	server.ListenAndServe()
}

func startPeriodicHealthCheck(m *MicroService, interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for t := range ticker.C {
		fmt.Println("Checking health of Service:", m.Route, " ---tick:", t)
		handleHealthCheck(m)
	}
}

func handleHealthCheck(m *MicroService) bool {
	change := false
	//var weightSum float32 = 0.0
	for i, inst := range m.Instances {
		_, err := dialer.Dial("tcp", inst.Host)
		if err != nil {
			if !m.isBlacklisted(i) {
				m.blackList(i, true)
				change = true
			}
		} else {
			if m.isBlacklisted(i) {
				m.blackList(i, false)
				change = true
			}
			//weightSum += inst.Weight
		}
	}
	//m.Mutex.Lock()
	//m.WeightSum = weightSum
	//m.Mutex.Unlock()
	return change
}

func NewPylonHandler(p *Pylon) http.HandlerFunc {
	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.getLoadBalancedInstance()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		fmt.Println("Instance is " + inst.Host)

		m.ReqCount <- 1
		inst.ReqCount <- 1

		fmt.Println("Serving new request, count: " + strconv.Itoa(len(m.ReqCount)))
		proxy := proxyPool.Get()
		SetUpProxy(proxy, m, inst.Host)
		proxy.ServeHTTP(w, r)
		proxyPool.Put(proxy)

		<-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) getLoadBalancedInstance() (*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")
	}

	instances := make([]*Instance, instCount)
	copy(instances, m.Instances)

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

		if err != nil {
			return nil, -1, err
		}

		if m.isBlacklisted(idx) {
			instances[idx] = nil
		} else {
			m.LastUsedIdx.Set(idx)
			return instances[idx], idx, nil
		}
	}
}

func nextRoundRobinInstIdx(instances []*Instance, idx int) (int, error) {
	tryCount := 1
	instCount := len(instances)
	lastNonNil := -1
	for {
		inst := instances[idx]
		if inst != nil {
			if inst.isRoundRobinPicked() {
				break
			}
			lastNonNil = idx
		}
		idx++
		tryCount++
		if tryCount > instCount {
			if lastNonNil != -1 {
				return lastNonNil, nil
			}
			return -1, errors.New("No instance can be round robin picked")
		}
		if idx >= instCount {
			idx = 0
		}
	}
	return idx, nil
}

func getLeastConInstIdx(instances []*Instance) int {
	minLoad := maxFloat32
	idx := 0

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

	return idx
}

func getRandomInstIdx(instances []*Instance) int {
	var weightSum float32 = 0.0
	for _, inst := range instances {
		if inst == nil {
			continue
		}
		weightSum += inst.Weight
	}

	r := rand.Float32() * weightSum

	for i, inst := range instances {
		if inst == nil {
			continue
		}
		r -= inst.Weight
		if r < 0 {
			return i
		}
	}

	return 0
}

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

func (m *MicroService) blackList(idx int, blacklist bool) {
	m.Mutex.Lock()
	if blacklist {
		m.BlackList[idx] = blacklisted
	} else {
		delete(m.BlackList, idx)
	}
	m.Mutex.Unlock()
}

func (m *MicroService) blackListHost(host string, blacklist bool) {
	for idx, inst := range m.Instances {
		if inst.Host == host {
			m.blackList(idx, blacklist)
		}
	}
}

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 mysqlclient

import (
	"strconv"

	"github.com/pubnative/mysqlproto-go"
)

type Rows struct {
	resultSet mysqlproto.ResultSet
	packet    []byte
	offset    uint64
	eof       bool
	err       error
}

func (r *Rows) Next() bool {
	if r.eof {
		return false
	}

	if r.err != nil {
		return false
	}

	packet, err := r.resultSet.Row()
	if err != nil {
		r.err = err
		r.eof = true
		return false
	}

	if packet == nil {
		r.eof = true
		return false
	} else {
		r.packet = packet
		r.offset = 0
		return true
	}
}

func (r *Rows) String() string {
	value, offset := mysqlproto.ParseString(r.packet, r.offset)
	r.offset = offset
	return value
}

func (r *Rows) Int() int {
	num, err := strconv.Atoi(r.String())
	if err != nil {
		r.err = err
	}

	return num
}

func (r *Rows) Int8() int8 {
	num, err := strconv.ParseInt(r.String(), 10, 8)
	if err != nil {
		r.err = err
	}
	return int8(num)
}

func (r *Rows) Int16() int16 {
	num, err := strconv.ParseInt(r.String(), 10, 16)
	if err != nil {
		r.err = err
	}
	return int16(num)
}

func (r *Rows) Int32() int32 {
	num, err := strconv.ParseInt(r.String(), 10, 32)
	if err != nil {
		r.err = err
	}
	return int32(num)
}

func (r *Rows) Int64() int64 {
	num, err := strconv.ParseInt(r.String(), 10, 64)
	if err != nil {
		r.err = err
	}
	return int64(num)
}

func (r *Rows) Float32() float32 {
	num, err := strconv.ParseFloat(r.String(), 32)
	if err != nil {
		r.err = err
	}
	return float32(num)
}

func (r *Rows) Float64() float64 {
	num, err := strconv.ParseFloat(r.String(), 64)
	if err != nil {
		r.err = err
	}
	return num
}

func (r *Rows) LastError() error {
	return r.err
}

func (c Conn) Query(sql string) (*Rows, error) {
	req := mysqlproto.ComQueryRequest([]byte(sql))
	if _, err := c.conn.Write(req); err != nil {
		return nil, err
	}

	resultSet, err := mysqlproto.ComQueryResponse(c.conn)
	if err != nil {
		return nil, err
	}

	return &Rows{resultSet: resultSet}, nil
}
<commit_msg>Allow to read raw bytes from table<commit_after>package mysqlclient

import (
	"strconv"

	"github.com/pubnative/mysqlproto-go"
)

type Rows struct {
	resultSet mysqlproto.ResultSet
	packet    []byte
	offset    uint64
	eof       bool
	err       error
}

func (r *Rows) Next() bool {
	if r.eof {
		return false
	}

	if r.err != nil {
		return false
	}

	packet, err := r.resultSet.Row()
	if err != nil {
		r.err = err
		r.eof = true
		return false
	}

	if packet == nil {
		r.eof = true
		return false
	} else {
		r.packet = packet
		r.offset = 0
		return true
	}
}

func (r *Rows) Bytes() []byte {
	value, offset := mysqlproto.ReadRowValue(r.packet, r.offset)
	r.offset = offset
	return value
}

func (r *Rows) String() string {
	return string(r.Bytes())
}

func (r *Rows) Int() int {
	num, err := strconv.Atoi(r.String())
	if err != nil {
		r.err = err
	}

	return num
}

func (r *Rows) Int8() int8 {
	num, err := strconv.ParseInt(r.String(), 10, 8)
	if err != nil {
		r.err = err
	}
	return int8(num)
}

func (r *Rows) Int16() int16 {
	num, err := strconv.ParseInt(r.String(), 10, 16)
	if err != nil {
		r.err = err
	}
	return int16(num)
}

func (r *Rows) Int32() int32 {
	num, err := strconv.ParseInt(r.String(), 10, 32)
	if err != nil {
		r.err = err
	}
	return int32(num)
}

func (r *Rows) Int64() int64 {
	num, err := strconv.ParseInt(r.String(), 10, 64)
	if err != nil {
		r.err = err
	}
	return int64(num)
}

func (r *Rows) Float32() float32 {
	num, err := strconv.ParseFloat(r.String(), 32)
	if err != nil {
		r.err = err
	}
	return float32(num)
}

func (r *Rows) Float64() float64 {
	num, err := strconv.ParseFloat(r.String(), 64)
	if err != nil {
		r.err = err
	}
	return num
}

func (r *Rows) LastError() error {
	return r.err
}

func (c Conn) Query(sql string) (*Rows, error) {
	req := mysqlproto.ComQueryRequest([]byte(sql))
	if _, err := c.conn.Write(req); err != nil {
		return nil, err
	}

	resultSet, err := mysqlproto.ComQueryResponse(c.conn)
	if err != nil {
		return nil, err
	}

	return &Rows{resultSet: resultSet}, nil
}
<|endoftext|>


<commit_before>package gorethink

import (
	"fmt"
	"strconv"
	"strings"

	p "github.com/dancannon/gorethink/ql2"
)

// A Query represents a query ready to be sent to the database, A Query differs
// from a Term as it contains both a query type and token. These values are used
// by the database to determine if the query is continuing a previous request
// and also allows the driver to identify the response as they can come out of
// order.
type Query struct {
	Type  p.Query_QueryType
	Token int64
	Term  *Term
	Opts  map[string]interface{}
}

func (q *Query) build() []interface{} {
	res := []interface{}{int(q.Type)}
	if q.Term != nil {
		res = append(res, q.Term.build())
	}

	if len(q.Opts) > 0 {
		res = append(res, q.Opts)
	}

	return res
}

type termsList []Term
type termsObj map[string]Term

// A Term represents a query that is being built. Terms consist of a an array of
// "sub-terms" and a term type. When a Term is a sub-term the first element of
// the terms data is its parent Term.
//
// When built the term becomes a JSON array, for more information on the format
// see http://rethinkdb.com/docs/writing-drivers/.
type Term struct {
	name     string
	rootTerm bool
	termType p.Term_TermType
	data     interface{}
	args     []Term
	optArgs  map[string]Term
}

// build takes the query tree and prepares it to be sent as a JSON
// expression
func (t Term) build() interface{} {
	switch t.termType {
	case p.Term_DATUM:
		return t.data
	case p.Term_MAKE_OBJ:
		res := map[string]interface{}{}
		for k, v := range t.optArgs {
			res[k] = v.build()
		}
		return res
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return map[string]interface{}{
				"$reql_type$": "BINARY",
				"data":        t.data,
			}
		}
	}

	args := []interface{}{}
	optArgs := map[string]interface{}{}

	for _, v := range t.args {
		args = append(args, v.build())
	}

	for k, v := range t.optArgs {
		optArgs[k] = v.build()
	}

	ret := []interface{}{int(t.termType)}

	if len(args) > 0 {
		ret = append(ret, args)
	}
	if len(optArgs) > 0 {
		ret = append(ret, optArgs)
	}

	return ret
}

// String returns a string representation of the query tree
func (t Term) String() string {
	switch t.termType {
	case p.Term_MAKE_ARRAY:
		return fmt.Sprintf("[%s]", strings.Join(argsToStringSlice(t.args), ", "))
	case p.Term_MAKE_OBJ:
		return fmt.Sprintf("{%s}", strings.Join(optArgsToStringSlice(t.optArgs), ", "))
	case p.Term_FUNC:
		// Get string representation of each argument
		args := []string{}
		for _, v := range t.args[0].args {
			args = append(args, fmt.Sprintf("var_%d", v.data))
		}

		return fmt.Sprintf("func(%s r.Term) r.Term { return %s }",
			strings.Join(args, ", "),
			t.args[1].String(),
		)
	case p.Term_VAR:
		return fmt.Sprintf("var_%s", t.args[0])
	case p.Term_IMPLICIT_VAR:
		return "r.Row"
	case p.Term_DATUM:
		switch v := t.data.(type) {
		case string:
			return strconv.Quote(v)
		default:
			return fmt.Sprintf("%v", v)
		}
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return fmt.Sprintf("r.binary(<data>)")
		}
	}

	if t.rootTerm {
		return fmt.Sprintf("r.%s(%s)", t.name, strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
	}
	return fmt.Sprintf("%s.%s(%s)", t.args[0].String(), t.name, strings.Join(allArgsToStringSlice(t.args[1:], t.optArgs), ", "))
}

// OptArgs is an interface used to represent a terms optional arguments. All
// optional argument types have a toMap function, the returned map can be encoded
// and sent as part of the query.
type OptArgs interface {
	toMap() map[string]interface{}
}

// WriteResponse is a helper type used when dealing with the response of a
// write query. It is also returned by the RunWrite function.
type WriteResponse struct {
	Errors        int              `gorethink:"errors"`
	Inserted      int              `gorethink:"inserted"`
	Updated       int              `gorethink:"updated"`
	Unchanged     int              `gorethink:"unchanged"`
	Replaced      int              `gorethink:"replaced"`
	Renamed       int              `gorethink:"renamed"`
	Skipped       int              `gorethink:"skipped"`
	Deleted       int              `gorethink:"deleted"`
	Created       int              `gorethink:"created"`
	DBsCreated    int              `gorethink:"dbs_created"`
	TablesCreated int              `gorethink:"tables_created"`
	Dropped       int              `gorethink:"dropped"`
	DBsDropped    int              `gorethink:"dbs_dropped"`
	TablesDropped int              `gorethink:"tables_dropped"`
	GeneratedKeys []string         `gorethink:"generated_keys"`
	FirstError    string           `gorethink:"first_error"` // populated if Errors > 0
	ConfigChanges []ChangeResponse `gorethink:"config_changes"`
	Changes       []ChangeResponse
}

// ChangeResponse is a helper type used when dealing with changefeeds. The type
// contains both the value before the query and the new value.
type ChangeResponse struct {
	NewValue interface{} `gorethink:"new_val"`
	OldValue interface{} `gorethink:"old_val"`
}

// RunOpts contains the optional arguments for the Run function.
type RunOpts struct {
	Db             interface{} `gorethink:"db,omitempty"`
	Profile        interface{} `gorethink:"profile,omitempty"`
	UseOutdated    interface{} `gorethink:"use_outdated,omitempty"`
	ArrayLimit     interface{} `gorethink:"array_limit,omitempty"`
	TimeFormat     interface{} `gorethink:"time_format,omitempty"`
	GroupFormat    interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat   interface{} `gorethink:"binary_format,omitempty"`
	GeometryFormat interface{} `gorethink:"geometry_format,omitempty"`

	MinBatchRows              interface{} `gorethink:"min_batch_rows,omitempty"`
	MaxBatchRows              interface{} `gorethink:"max_batch_rows,omitempty"`
	MaxBatchBytes             interface{} `gorethink:"max_batch_bytes,omitempty"`
	MaxBatchSeconds           interface{} `gorethink:"max_batch_seconds,omitempty"`
	FirstBatchScaledownFactor interface{} `gorethink:"first_batch_scaledown_factor,omitempty"`
}

func (o *RunOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Run runs a query using the given connection.
//
//	rows, err := query.Run(sess)
//	if err != nil {
//		// error
//	}
//
//  var doc MyDocumentType
//	for rows.Next(&doc) {
//      // Do something with document
//	}
func (t Term) Run(s *Session, optArgs ...RunOpts) (*Cursor, error) {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}

	return s.Query(s.newQuery(t, opts))
}

// RunWrite runs a query using the given connection but unlike Run automatically
// scans the result into a variable of type WriteResponss. This function should be used
// if you are running a write query (such as Insert,  Update, TableCreate, etc...)
//
//	res, err := r.Db("database").Table("table").Insert(doc).RunWrite(sess)
func (t Term) RunWrite(s *Session, optArgs ...RunOpts) (WriteResponse, error) {
	var response WriteResponse

	res, err := t.Run(s, optArgs...)
	if err != nil {
		return response, err
	}

	if err = res.One(&response); err != nil {
		return response, err
	}

	if err = res.Close(); err != nil {
		return response, err
	}

	return response, nil
}

// ExecOpts contains the optional arguments for the Exec function and  inherits
// its options from RunOpts, the only difference is the addition of the NoReply
// field.
//
// When NoReply is true it causes the driver not to wait to receive the result
// and return immediately.
type ExecOpts struct {
	Db             interface{} `gorethink:"db,omitempty"`
	Profile        interface{} `gorethink:"profile,omitempty"`
	UseOutdated    interface{} `gorethink:"use_outdated,omitempty"`
	ArrayLimit     interface{} `gorethink:"array_limit,omitempty"`
	TimeFormat     interface{} `gorethink:"time_format,omitempty"`
	GroupFormat    interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat   interface{} `gorethink:"binary_format,omitempty"`
	GeometryFormat interface{} `gorethink:"geometry_format,omitempty"`

	MinBatchRows              interface{} `gorethink:"min_batch_rows,omitempty"`
	MaxBatchRows              interface{} `gorethink:"max_batch_rows,omitempty"`
	MaxBatchBytes             interface{} `gorethink:"max_batch_bytes,omitempty"`
	MaxBatchSeconds           interface{} `gorethink:"max_batch_seconds,omitempty"`
	FirstBatchScaledownFactor interface{} `gorethink:"first_batch_scaledown_factor,omitempty"`

	NoReply interface{} `gorethink:"noreply,omitempty"`
}

func (o *ExecOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Exec runs the query but does not return the result. Exec will still wait for
// the response to be received unless the NoReply field is true.
//
//	err := r.Db("database").Table("table").Insert(doc).Exec(sess, r.ExecOpts{
//		NoReply: true,
//	})
func (t Term) Exec(s *Session, optArgs ...ExecOpts) error {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}

	return s.Exec(s.newQuery(t, opts))
}
<commit_msg>Return an error if any write errors are returned when calling RunWrite<commit_after>package gorethink

import (
	"fmt"
	"strconv"
	"strings"

	p "github.com/dancannon/gorethink/ql2"
)

// A Query represents a query ready to be sent to the database, A Query differs
// from a Term as it contains both a query type and token. These values are used
// by the database to determine if the query is continuing a previous request
// and also allows the driver to identify the response as they can come out of
// order.
type Query struct {
	Type  p.Query_QueryType
	Token int64
	Term  *Term
	Opts  map[string]interface{}
}

func (q *Query) build() []interface{} {
	res := []interface{}{int(q.Type)}
	if q.Term != nil {
		res = append(res, q.Term.build())
	}

	if len(q.Opts) > 0 {
		res = append(res, q.Opts)
	}

	return res
}

type termsList []Term
type termsObj map[string]Term

// A Term represents a query that is being built. Terms consist of a an array of
// "sub-terms" and a term type. When a Term is a sub-term the first element of
// the terms data is its parent Term.
//
// When built the term becomes a JSON array, for more information on the format
// see http://rethinkdb.com/docs/writing-drivers/.
type Term struct {
	name     string
	rootTerm bool
	termType p.Term_TermType
	data     interface{}
	args     []Term
	optArgs  map[string]Term
}

// build takes the query tree and prepares it to be sent as a JSON
// expression
func (t Term) build() interface{} {
	switch t.termType {
	case p.Term_DATUM:
		return t.data
	case p.Term_MAKE_OBJ:
		res := map[string]interface{}{}
		for k, v := range t.optArgs {
			res[k] = v.build()
		}
		return res
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return map[string]interface{}{
				"$reql_type$": "BINARY",
				"data":        t.data,
			}
		}
	}

	args := []interface{}{}
	optArgs := map[string]interface{}{}

	for _, v := range t.args {
		args = append(args, v.build())
	}

	for k, v := range t.optArgs {
		optArgs[k] = v.build()
	}

	ret := []interface{}{int(t.termType)}

	if len(args) > 0 {
		ret = append(ret, args)
	}
	if len(optArgs) > 0 {
		ret = append(ret, optArgs)
	}

	return ret
}

// String returns a string representation of the query tree
func (t Term) String() string {
	switch t.termType {
	case p.Term_MAKE_ARRAY:
		return fmt.Sprintf("[%s]", strings.Join(argsToStringSlice(t.args), ", "))
	case p.Term_MAKE_OBJ:
		return fmt.Sprintf("{%s}", strings.Join(optArgsToStringSlice(t.optArgs), ", "))
	case p.Term_FUNC:
		// Get string representation of each argument
		args := []string{}
		for _, v := range t.args[0].args {
			args = append(args, fmt.Sprintf("var_%d", v.data))
		}

		return fmt.Sprintf("func(%s r.Term) r.Term { return %s }",
			strings.Join(args, ", "),
			t.args[1].String(),
		)
	case p.Term_VAR:
		return fmt.Sprintf("var_%s", t.args[0])
	case p.Term_IMPLICIT_VAR:
		return "r.Row"
	case p.Term_DATUM:
		switch v := t.data.(type) {
		case string:
			return strconv.Quote(v)
		default:
			return fmt.Sprintf("%v", v)
		}
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return fmt.Sprintf("r.binary(<data>)")
		}
	}

	if t.rootTerm {
		return fmt.Sprintf("r.%s(%s)", t.name, strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
	}
	return fmt.Sprintf("%s.%s(%s)", t.args[0].String(), t.name, strings.Join(allArgsToStringSlice(t.args[1:], t.optArgs), ", "))
}

// OptArgs is an interface used to represent a terms optional arguments. All
// optional argument types have a toMap function, the returned map can be encoded
// and sent as part of the query.
type OptArgs interface {
	toMap() map[string]interface{}
}

// WriteResponse is a helper type used when dealing with the response of a
// write query. It is also returned by the RunWrite function.
type WriteResponse struct {
	Errors        int              `gorethink:"errors"`
	Inserted      int              `gorethink:"inserted"`
	Updated       int              `gorethink:"updated"`
	Unchanged     int              `gorethink:"unchanged"`
	Replaced      int              `gorethink:"replaced"`
	Renamed       int              `gorethink:"renamed"`
	Skipped       int              `gorethink:"skipped"`
	Deleted       int              `gorethink:"deleted"`
	Created       int              `gorethink:"created"`
	DBsCreated    int              `gorethink:"dbs_created"`
	TablesCreated int              `gorethink:"tables_created"`
	Dropped       int              `gorethink:"dropped"`
	DBsDropped    int              `gorethink:"dbs_dropped"`
	TablesDropped int              `gorethink:"tables_dropped"`
	GeneratedKeys []string         `gorethink:"generated_keys"`
	FirstError    string           `gorethink:"first_error"` // populated if Errors > 0
	ConfigChanges []ChangeResponse `gorethink:"config_changes"`
	Changes       []ChangeResponse
}

// ChangeResponse is a helper type used when dealing with changefeeds. The type
// contains both the value before the query and the new value.
type ChangeResponse struct {
	NewValue interface{} `gorethink:"new_val"`
	OldValue interface{} `gorethink:"old_val"`
}

// RunOpts contains the optional arguments for the Run function.
type RunOpts struct {
	Db             interface{} `gorethink:"db,omitempty"`
	Profile        interface{} `gorethink:"profile,omitempty"`
	UseOutdated    interface{} `gorethink:"use_outdated,omitempty"`
	ArrayLimit     interface{} `gorethink:"array_limit,omitempty"`
	TimeFormat     interface{} `gorethink:"time_format,omitempty"`
	GroupFormat    interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat   interface{} `gorethink:"binary_format,omitempty"`
	GeometryFormat interface{} `gorethink:"geometry_format,omitempty"`

	MinBatchRows              interface{} `gorethink:"min_batch_rows,omitempty"`
	MaxBatchRows              interface{} `gorethink:"max_batch_rows,omitempty"`
	MaxBatchBytes             interface{} `gorethink:"max_batch_bytes,omitempty"`
	MaxBatchSeconds           interface{} `gorethink:"max_batch_seconds,omitempty"`
	FirstBatchScaledownFactor interface{} `gorethink:"first_batch_scaledown_factor,omitempty"`
}

func (o *RunOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Run runs a query using the given connection.
//
//	rows, err := query.Run(sess)
//	if err != nil {
//		// error
//	}
//
//  var doc MyDocumentType
//	for rows.Next(&doc) {
//      // Do something with document
//	}
func (t Term) Run(s *Session, optArgs ...RunOpts) (*Cursor, error) {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}

	return s.Query(s.newQuery(t, opts))
}

// RunWrite runs a query using the given connection but unlike Run automatically
// scans the result into a variable of type WriteResponss. This function should be used
// if you are running a write query (such as Insert,  Update, TableCreate, etc...).
//
// If an error occurs when running the write query the first error is returned.
//
//	res, err := r.Db("database").Table("table").Insert(doc).RunWrite(sess)
func (t Term) RunWrite(s *Session, optArgs ...RunOpts) (WriteResponse, error) {
	var response WriteResponse

	res, err := t.Run(s, optArgs...)
	if err != nil {
		return response, err
	}

	if err = res.One(&response); err != nil {
		return response, err
	}

	if err = res.Close(); err != nil {
		return response, err
	}

	if response.Errors > 0 {
		return response, fmt.Errorf(response.FirstError)
	}

	return response, nil
}

// ExecOpts contains the optional arguments for the Exec function and  inherits
// its options from RunOpts, the only difference is the addition of the NoReply
// field.
//
// When NoReply is true it causes the driver not to wait to receive the result
// and return immediately.
type ExecOpts struct {
	Db             interface{} `gorethink:"db,omitempty"`
	Profile        interface{} `gorethink:"profile,omitempty"`
	UseOutdated    interface{} `gorethink:"use_outdated,omitempty"`
	ArrayLimit     interface{} `gorethink:"array_limit,omitempty"`
	TimeFormat     interface{} `gorethink:"time_format,omitempty"`
	GroupFormat    interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat   interface{} `gorethink:"binary_format,omitempty"`
	GeometryFormat interface{} `gorethink:"geometry_format,omitempty"`

	MinBatchRows              interface{} `gorethink:"min_batch_rows,omitempty"`
	MaxBatchRows              interface{} `gorethink:"max_batch_rows,omitempty"`
	MaxBatchBytes             interface{} `gorethink:"max_batch_bytes,omitempty"`
	MaxBatchSeconds           interface{} `gorethink:"max_batch_seconds,omitempty"`
	FirstBatchScaledownFactor interface{} `gorethink:"first_batch_scaledown_factor,omitempty"`

	NoReply interface{} `gorethink:"noreply,omitempty"`
}

func (o *ExecOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Exec runs the query but does not return the result. Exec will still wait for
// the response to be received unless the NoReply field is true.
//
//	err := r.Db("database").Table("table").Insert(doc).Exec(sess, r.ExecOpts{
//		NoReply: true,
//	})
func (t Term) Exec(s *Session, optArgs ...ExecOpts) error {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}

	return s.Exec(s.newQuery(t, opts))
}
<|endoftext|>


<commit_before>package presto

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

const (
	initialRetry            = 50 * time.Millisecond
	maxRetry                = 800 * time.Millisecond
	ProgressUnknown float64 = -1.0
)

// Query represents an open query to presto.
type Query struct {
	host    string
	user    string
	source  string
	catalog string
	schema  string

	query   string
	closed  bool
	id      string
	columns []string

	bufferedRows [][]interface{}
	state        string
	progress     float64
	nextUri      string
}

type queryResult struct {
	ID               string          `json:"id"`
	InfoUri          string          `json:"infoUri"`
	NextUri          string          `json:"nextUri"`
	PartialCancelUri string          `json:"PartialCancelUri"`
	Data             [][]interface{} `json:"data"`
	Columns          []struct {
		Name string `json:"name"`
	} `json:"columns"`
	Error struct {
		ErrorCode   int `json:"errorCode"`
		FailureInfo struct {
			Message string `json:"message"`
		} `json:"failureInfo"`
	} `json:"error"`
	Stats struct {
		State           string `json:"state"`
		Scheduled       bool   `json:"scheduled"`
		CompletedSplits int    `json:"completedSplits"`
		TotalSplits     int    `json:"totalSplits"`
	} `json:"stats"`
}

func NewQuery(host, user, source, catalog, schema, query string) (*Query, error) {
	if user == "" {
		user = "anonymous"
	}

	if source == "" {
		source = userAgent
	}

	if catalog == "" {
		catalog = "default"
	}

	if schema == "" {
		schema = "default"
	}

	q := &Query{
		host:    host,
		user:    user,
		catalog: catalog,
		schema:  schema,
		query:   query,
	}

	err := q.postQuery()
	if err != nil {
		return nil, err
	}

	// Immediately fetch one result set, to fill the column names.
	err = q.fetchNext()
	if err != nil {
		return nil, err
	}

	return q, nil
}

// Columns returns a list of the column names for the query.
func (q *Query) Columns() []string {
	return q.columns
}

// Progress returns the current progress of the query, which is the proportion
// of splits completed. Progress is only updated on calls to Next.
//
// If the total number of splits is not determined, Progress returns
// ProgressUnknown.
func (q *Query) Progress() float64 {
	return q.progress
}

// Close closes the query, and cancels it if started.
func (q *Query) Close() error {
	if q.closed {
		return nil
	}

	q.closed = true
	req, _ := http.NewRequest("DELETE", q.nextUri, nil)
	resp, err := q.makeRequest(req)
	if err != nil {
		return err
	}

	if resp.StatusCode != 204 {
		return fmt.Errorf("unexpected http status: %s", resp.Status)
	}

	return nil
}

// Next retrieves the next row from the dataset, fetching more if need be.
func (q *Query) Next() ([]interface{}, error) {
	retry := initialRetry
	for !q.closed && len(q.bufferedRows) == 0 {
		err := q.fetchNext()
		if err != nil {
			return nil, err
		}

		if len(q.bufferedRows) == 0 {
			time.Sleep(retry)
			retry *= 2
			if retry > maxRetry {
				retry = maxRetry
			}
		}
	}

	if len(q.bufferedRows) > 0 {
		row := q.bufferedRows[0]
		q.bufferedRows = q.bufferedRows[1:]
		return row, nil
	} else {
		return nil, nil
	}
}

func (q *Query) postQuery() error {
	queryUrl := fmt.Sprintf("%s/v1/statement", q.host)
	req, _ := http.NewRequest("POST", queryUrl, strings.NewReader(q.query))
	result, err := q.fetchResult(req)
	if err != nil {
		q.closed = true
		return err
	}

	q.id = result.ID
	q.nextUri = result.NextUri
	return nil
}

func (q *Query) fetchNext() error {
	req, _ := http.NewRequest("GET", q.nextUri, nil)
	result, err := q.fetchResult(req)
	if err != nil {
		q.closed = true
		return err
	}

	q.bufferedRows = result.Data
	q.state = result.Stats.State

	if q.columns == nil {
		q.columns = make([]string, len(result.Columns))
		for i, col := range result.Columns {
			q.columns[i] = col.Name
		}
	}

	if result.Stats.Scheduled {
		q.progress = float64(result.Stats.CompletedSplits) / float64(result.Stats.TotalSplits)
	} else {
		q.progress = ProgressUnknown
	}

	q.nextUri = result.NextUri
	if result.NextUri == "" {
		q.closed = true
	}

	return nil
}

func (q *Query) fetchResult(req *http.Request) (*queryResult, error) {
	resp, err := q.makeRequest(req)
	if err != nil {
		return nil, err
	}

	result, err := q.readResult(resp)
	if err != nil {
		return result, err
	}

	if result.Error.FailureInfo.Message != "" {
		return result, fmt.Errorf("query failed: %s", result.Error.FailureInfo.Message)
	}

	return result, nil
}

func (q *Query) makeRequest(req *http.Request) (*http.Response, error) {
	req.Header.Add("User-Agent", userAgent)
	req.Header.Add("X-Presto-User", q.user)
	req.Header.Add("X-Presto-Catalog", q.catalog)
	req.Header.Add("X-Presto-Schema", q.schema)

	// Sometimes presto returns a 503 to indicate that results aren't yet
	// available, and we should retry after waiting a bit.
	retry := initialRetry
	for {
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return nil, err
		}

		if resp.StatusCode == 200 {
			return resp, nil
		} else if resp.StatusCode != 503 {
			return nil, fmt.Errorf("unexpected http status: %s", resp.Status)
		}

		time.Sleep(retry)
		retry *= 2
		if retry > maxRetry {
			retry = maxRetry
		}
	}
}

func (q *Query) readResult(resp *http.Response) (*queryResult, error) {
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	result := queryResult{}
	err = json.Unmarshal(body, &result)
	if err != nil {
		return nil, fmt.Errorf("error decoding json response from presto: %s", err)
	}

	return &result, nil
}
<commit_msg>Added support for X-Presto-Source<commit_after>package presto

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

const (
	initialRetry            = 50 * time.Millisecond
	maxRetry                = 800 * time.Millisecond
	ProgressUnknown float64 = -1.0
)

// Query represents an open query to presto.
type Query struct {
	host    string
	user    string
	source  string
	catalog string
	schema  string

	query   string
	closed  bool
	id      string
	columns []string

	bufferedRows [][]interface{}
	state        string
	progress     float64
	nextUri      string
}

type queryResult struct {
	ID               string          `json:"id"`
	InfoUri          string          `json:"infoUri"`
	NextUri          string          `json:"nextUri"`
	PartialCancelUri string          `json:"PartialCancelUri"`
	Data             [][]interface{} `json:"data"`
	Columns          []struct {
		Name string `json:"name"`
	} `json:"columns"`
	Error struct {
		ErrorCode   int `json:"errorCode"`
		FailureInfo struct {
			Message string `json:"message"`
		} `json:"failureInfo"`
	} `json:"error"`
	Stats struct {
		State           string `json:"state"`
		Scheduled       bool   `json:"scheduled"`
		CompletedSplits int    `json:"completedSplits"`
		TotalSplits     int    `json:"totalSplits"`
	} `json:"stats"`
}

func NewQuery(host, user, source, catalog, schema, query string) (*Query, error) {
	if user == "" {
		user = "anonymous"
	}

	if source == "" {
		source = userAgent
	}

	if catalog == "" {
		catalog = "default"
	}

	if schema == "" {
		schema = "default"
	}

	q := &Query{
		host:    host,
		user:    user,
		source:  source,
		catalog: catalog,
		schema:  schema,
		query:   query,
	}

	err := q.postQuery()
	if err != nil {
		return nil, err
	}

	// Immediately fetch one result set, to fill the column names.
	err = q.fetchNext()
	if err != nil {
		return nil, err
	}

	return q, nil
}

// Columns returns a list of the column names for the query.
func (q *Query) Columns() []string {
	return q.columns
}

// Progress returns the current progress of the query, which is the proportion
// of splits completed. Progress is only updated on calls to Next.
//
// If the total number of splits is not determined, Progress returns
// ProgressUnknown.
func (q *Query) Progress() float64 {
	return q.progress
}

// Close closes the query, and cancels it if started.
func (q *Query) Close() error {
	if q.closed {
		return nil
	}

	q.closed = true
	req, _ := http.NewRequest("DELETE", q.nextUri, nil)
	resp, err := q.makeRequest(req)
	if err != nil {
		return err
	}

	if resp.StatusCode != 204 {
		return fmt.Errorf("unexpected http status: %s", resp.Status)
	}

	return nil
}

// Next retrieves the next row from the dataset, fetching more if need be.
func (q *Query) Next() ([]interface{}, error) {
	retry := initialRetry
	for !q.closed && len(q.bufferedRows) == 0 {
		err := q.fetchNext()
		if err != nil {
			return nil, err
		}

		if len(q.bufferedRows) == 0 {
			time.Sleep(retry)
			retry *= 2
			if retry > maxRetry {
				retry = maxRetry
			}
		}
	}

	if len(q.bufferedRows) > 0 {
		row := q.bufferedRows[0]
		q.bufferedRows = q.bufferedRows[1:]
		return row, nil
	} else {
		return nil, nil
	}
}

func (q *Query) postQuery() error {
	queryUrl := fmt.Sprintf("%s/v1/statement", q.host)
	req, _ := http.NewRequest("POST", queryUrl, strings.NewReader(q.query))
	result, err := q.fetchResult(req)
	if err != nil {
		q.closed = true
		return err
	}

	q.id = result.ID
	q.nextUri = result.NextUri
	return nil
}

func (q *Query) fetchNext() error {
	req, _ := http.NewRequest("GET", q.nextUri, nil)
	result, err := q.fetchResult(req)
	if err != nil {
		q.closed = true
		return err
	}

	q.bufferedRows = result.Data
	q.state = result.Stats.State

	if q.columns == nil {
		q.columns = make([]string, len(result.Columns))
		for i, col := range result.Columns {
			q.columns[i] = col.Name
		}
	}

	if result.Stats.Scheduled {
		q.progress = float64(result.Stats.CompletedSplits) / float64(result.Stats.TotalSplits)
	} else {
		q.progress = ProgressUnknown
	}

	q.nextUri = result.NextUri
	if result.NextUri == "" {
		q.closed = true
	}

	return nil
}

func (q *Query) fetchResult(req *http.Request) (*queryResult, error) {
	resp, err := q.makeRequest(req)
	if err != nil {
		return nil, err
	}

	result, err := q.readResult(resp)
	if err != nil {
		return result, err
	}

	if result.Error.FailureInfo.Message != "" {
		return result, fmt.Errorf("query failed: %s", result.Error.FailureInfo.Message)
	}

	return result, nil
}

func (q *Query) makeRequest(req *http.Request) (*http.Response, error) {
	req.Header.Add("User-Agent", userAgent)
	req.Header.Add("X-Presto-User", q.user)
	req.Header.Add("X-Presto-Catalog", q.catalog)
	req.Header.Add("X-Presto-Schema", q.schema)
	req.Header.Add("X-Presto-Source", q.source)

	// Sometimes presto returns a 503 to indicate that results aren't yet
	// available, and we should retry after waiting a bit.
	retry := initialRetry
	for {
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return nil, err
		}

		if resp.StatusCode == 200 {
			return resp, nil
		} else if resp.StatusCode != 503 {
			return nil, fmt.Errorf("unexpected http status: %s", resp.Status)
		}

		time.Sleep(retry)
		retry *= 2
		if retry > maxRetry {
			retry = maxRetry
		}
	}
}

func (q *Query) readResult(resp *http.Response) (*queryResult, error) {
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	result := queryResult{}
	err = json.Unmarshal(body, &result)
	if err != nil {
		return nil, fmt.Errorf("error decoding json response from presto: %s", err)
	}

	return &result, nil
}
<|endoftext|>


<commit_before>package gorethink

import (
	"fmt"
	"strconv"
	"strings"

	p "github.com/dancannon/gorethink/ql2"
)

type Query struct {
	Type  p.Query_QueryType
	Token int64
	Term  *Term
	Opts  map[string]interface{}
}

func (q *Query) build() []interface{} {
	res := []interface{}{int(q.Type)}
	if q.Term != nil {
		res = append(res, q.Term.build())
	}

	if len(q.Opts) > 0 {
		res = append(res, q.Opts)
	}

	return res
}

type termsList []Term
type termsObj map[string]Term
type Term struct {
	name     string
	rootTerm bool
	termType p.Term_TermType
	data     interface{}
	args     []Term
	optArgs  map[string]Term
}

// build takes the query tree and prepares it to be sent as a JSON
// expression
func (t Term) build() interface{} {
	switch t.termType {
	case p.Term_DATUM:
		return t.data
	case p.Term_MAKE_OBJ:
		res := map[string]interface{}{}
		for k, v := range t.optArgs {
			res[k] = v.build()
		}
		return res
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return map[string]interface{}{
				"$reql_type$": "BINARY",
				"data":        t.data,
			}
		}
	}

	args := []interface{}{}
	optArgs := map[string]interface{}{}

	for _, v := range t.args {
		args = append(args, v.build())
	}

	for k, v := range t.optArgs {
		optArgs[k] = v.build()
	}

	ret := []interface{}{int(t.termType)}

	if len(args) > 0 {
		ret = append(ret, args)
	}
	if len(optArgs) > 0 {
		ret = append(ret, optArgs)
	}

	return ret
}

// String returns a string representation of the query tree
func (t Term) String() string {
	switch t.termType {
	case p.Term_MAKE_ARRAY:
		return fmt.Sprintf("[%s]", strings.Join(argsToStringSlice(t.args), ", "))
	case p.Term_MAKE_OBJ:
		return fmt.Sprintf("{%s}", strings.Join(optArgsToStringSlice(t.optArgs), ", "))
	case p.Term_FUNC:
		// Get string representation of each argument
		args := []string{}
		for _, v := range t.args[0].args {
			args = append(args, fmt.Sprintf("var_%d", v.data))
		}

		return fmt.Sprintf("func(%s r.Term) r.Term { return %s }",
			strings.Join(args, ", "),
			t.args[1].String(),
		)
	case p.Term_VAR:
		return fmt.Sprintf("var_%s", t.args[0])
	case p.Term_IMPLICIT_VAR:
		return "r.Row"
	case p.Term_DATUM:
		switch v := t.data.(type) {
		case string:
			return strconv.Quote(v)
		default:
			return fmt.Sprintf("%v", v)
		}
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return fmt.Sprintf("r.binary(<data>)")
		}
	}

	if t.rootTerm {
		return fmt.Sprintf("r.%s(%s)", t.name, strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
	}
	return fmt.Sprintf("%s.%s(%s)", t.args[0].String(), t.name, strings.Join(allArgsToStringSlice(t.args[1:], t.optArgs), ", "))
}

type OptArgs interface {
	toMap() map[string]interface{}
}

type WriteResponse struct {
	Errors        int            `gorethink:"errors"`
	Inserted      int            `gorethink:"inserted"`
	Updated       int            `gorethink:"updadte"`
	Unchanged     int            `gorethink:"unchanged"`
	Replaced      int            `gorethink:"replaced"`
	Renamed       int            `gorethink:"renamed"`
	Skipped       int            `gorethink:"skipped"`
	Deleted       int            `gorethink:"deleted"`
	Created       int            `gorethink:"created"`
	DBsCreated    int            `gorethink:"dbs_created"`
	TablesCreated int            `gorethink:"tables_created"`
	Dropped       int            `gorethink:"dropped"`
	DBsDropped    int            `gorethink:"dbs_dropped"`
	TablesDropped int            `gorethink:"tables_dropped"`
	GeneratedKeys []string       `gorethink:"generated_keys"`
	FirstError    string         `gorethink:"first_error"` // populated if Errors > 0
	ConfigChanges []WriteChanges `gorethink:"config_changes"`
	Changes       []WriteChanges
}

type WriteChanges struct {
	NewValue interface{} `gorethink:"new_val"`
	OldValue interface{} `gorethink:"old_val"`
}

type RunOpts struct {
	Db             interface{} `gorethink:"db,omitempty"`
	Profile        interface{} `gorethink:"profile,omitempty"`
	UseOutdated    interface{} `gorethink:"use_outdated,omitempty"`
	ArrayLimit     interface{} `gorethink:"array_limit,omitempty"`
	TimeFormat     interface{} `gorethink:"time_format,omitempty"`
	GroupFormat    interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat   interface{} `gorethink:"binary_format,omitempty"`
	GeometryFormat interface{} `gorethink:"geometry_format,omitempty"`

	MinBatchRows              interface{} `gorethink:"min_batch_rows,omitempty"`
	MaxBatchRows              interface{} `gorethink:"max_batch_rows,omitempty"`
	MaxBatchBytes             interface{} `gorethink:"max_batch_bytes,omitempty"`
	MaxBatchSeconds           interface{} `gorethink:"max_batch_seconds,omitempty"`
	FirstBatchScaledownFactor interface{} `gorethink:"first_batch_scaledown_factor,omitempty"`
}

func (o *RunOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Run runs a query using the given connection.
//
//	rows, err := query.Run(sess)
//	if err != nil {
//		// error
//	}
//
//  var doc MyDocumentType
//	for rows.Next(&doc) {
//      // Do something with document
//	}
func (t Term) Run(s *Session, optArgs ...RunOpts) (*Cursor, error) {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}

	return s.Query(s.newQuery(t, opts))
}

// RunWrite runs a query using the given connection but unlike Run automatically
// scans the result into a variable of type WriteResponss. This function should be used
// if you are running a write query (such as Insert,  Update, TableCreate, etc...)
//
//	res, err := r.Db("database").Table("table").Insert(doc).RunWrite(sess)
func (t Term) RunWrite(s *Session, optArgs ...RunOpts) (WriteResponse, error) {
	var response WriteResponse
	res, err := t.Run(s, optArgs...)
	if err == nil {
		err = res.One(&response)
	}
	return response, err
}

// ExecOpts inherits its options from RunOpts, the only difference is the
// addition of the NoReply field.
//
// When NoReply is true it causes the driver not to wait to receive the result
// and return immediately.
type ExecOpts struct {
	Db             interface{} `gorethink:"db,omitempty"`
	Profile        interface{} `gorethink:"profile,omitempty"`
	UseOutdated    interface{} `gorethink:"use_outdated,omitempty"`
	ArrayLimit     interface{} `gorethink:"array_limit,omitempty"`
	TimeFormat     interface{} `gorethink:"time_format,omitempty"`
	GroupFormat    interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat   interface{} `gorethink:"binary_format,omitempty"`
	GeometryFormat interface{} `gorethink:"geometry_format,omitempty"`

	MinBatchRows              interface{} `gorethink:"min_batch_rows,omitempty"`
	MaxBatchRows              interface{} `gorethink:"max_batch_rows,omitempty"`
	MaxBatchBytes             interface{} `gorethink:"max_batch_bytes,omitempty"`
	MaxBatchSeconds           interface{} `gorethink:"max_batch_seconds,omitempty"`
	FirstBatchScaledownFactor interface{} `gorethink:"first_batch_scaledown_factor,omitempty"`

	NoReply interface{} `gorethink:"noreply,omitempty"`
}

func (o *ExecOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Exec runs the query but does not return the result. Exec will still wait for
// the response to be received unless the NoReply field is true.
//
//	res, err := r.Db("database").Table("table").Insert(doc).Exec(sess, r.ExecOpts{
//		NoReply: true,
//	})
func (t Term) Exec(s *Session, optArgs ...ExecOpts) error {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}

	return s.Exec(s.newQuery(t, opts))
}
<commit_msg>close *Cursor on defer<commit_after>package gorethink

import (
	"fmt"
	"strconv"
	"strings"

	p "github.com/dancannon/gorethink/ql2"
)

type Query struct {
	Type  p.Query_QueryType
	Token int64
	Term  *Term
	Opts  map[string]interface{}
}

func (q *Query) build() []interface{} {
	res := []interface{}{int(q.Type)}
	if q.Term != nil {
		res = append(res, q.Term.build())
	}

	if len(q.Opts) > 0 {
		res = append(res, q.Opts)
	}

	return res
}

type termsList []Term
type termsObj map[string]Term
type Term struct {
	name     string
	rootTerm bool
	termType p.Term_TermType
	data     interface{}
	args     []Term
	optArgs  map[string]Term
}

// build takes the query tree and prepares it to be sent as a JSON
// expression
func (t Term) build() interface{} {
	switch t.termType {
	case p.Term_DATUM:
		return t.data
	case p.Term_MAKE_OBJ:
		res := map[string]interface{}{}
		for k, v := range t.optArgs {
			res[k] = v.build()
		}
		return res
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return map[string]interface{}{
				"$reql_type$": "BINARY",
				"data":        t.data,
			}
		}
	}

	args := []interface{}{}
	optArgs := map[string]interface{}{}

	for _, v := range t.args {
		args = append(args, v.build())
	}

	for k, v := range t.optArgs {
		optArgs[k] = v.build()
	}

	ret := []interface{}{int(t.termType)}

	if len(args) > 0 {
		ret = append(ret, args)
	}
	if len(optArgs) > 0 {
		ret = append(ret, optArgs)
	}

	return ret
}

// String returns a string representation of the query tree
func (t Term) String() string {
	switch t.termType {
	case p.Term_MAKE_ARRAY:
		return fmt.Sprintf("[%s]", strings.Join(argsToStringSlice(t.args), ", "))
	case p.Term_MAKE_OBJ:
		return fmt.Sprintf("{%s}", strings.Join(optArgsToStringSlice(t.optArgs), ", "))
	case p.Term_FUNC:
		// Get string representation of each argument
		args := []string{}
		for _, v := range t.args[0].args {
			args = append(args, fmt.Sprintf("var_%d", v.data))
		}

		return fmt.Sprintf("func(%s r.Term) r.Term { return %s }",
			strings.Join(args, ", "),
			t.args[1].String(),
		)
	case p.Term_VAR:
		return fmt.Sprintf("var_%s", t.args[0])
	case p.Term_IMPLICIT_VAR:
		return "r.Row"
	case p.Term_DATUM:
		switch v := t.data.(type) {
		case string:
			return strconv.Quote(v)
		default:
			return fmt.Sprintf("%v", v)
		}
	case p.Term_BINARY:
		if len(t.args) == 0 {
			return fmt.Sprintf("r.binary(<data>)")
		}
	}

	if t.rootTerm {
		return fmt.Sprintf("r.%s(%s)", t.name, strings.Join(allArgsToStringSlice(t.args, t.optArgs), ", "))
	}
	return fmt.Sprintf("%s.%s(%s)", t.args[0].String(), t.name, strings.Join(allArgsToStringSlice(t.args[1:], t.optArgs), ", "))
}

type OptArgs interface {
	toMap() map[string]interface{}
}

type WriteResponse struct {
	Errors        int            `gorethink:"errors"`
	Inserted      int            `gorethink:"inserted"`
	Updated       int            `gorethink:"updadte"`
	Unchanged     int            `gorethink:"unchanged"`
	Replaced      int            `gorethink:"replaced"`
	Renamed       int            `gorethink:"renamed"`
	Skipped       int            `gorethink:"skipped"`
	Deleted       int            `gorethink:"deleted"`
	Created       int            `gorethink:"created"`
	DBsCreated    int            `gorethink:"dbs_created"`
	TablesCreated int            `gorethink:"tables_created"`
	Dropped       int            `gorethink:"dropped"`
	DBsDropped    int            `gorethink:"dbs_dropped"`
	TablesDropped int            `gorethink:"tables_dropped"`
	GeneratedKeys []string       `gorethink:"generated_keys"`
	FirstError    string         `gorethink:"first_error"` // populated if Errors > 0
	ConfigChanges []WriteChanges `gorethink:"config_changes"`
	Changes       []WriteChanges
}

type WriteChanges struct {
	NewValue interface{} `gorethink:"new_val"`
	OldValue interface{} `gorethink:"old_val"`
}

type RunOpts struct {
	Db             interface{} `gorethink:"db,omitempty"`
	Profile        interface{} `gorethink:"profile,omitempty"`
	UseOutdated    interface{} `gorethink:"use_outdated,omitempty"`
	ArrayLimit     interface{} `gorethink:"array_limit,omitempty"`
	TimeFormat     interface{} `gorethink:"time_format,omitempty"`
	GroupFormat    interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat   interface{} `gorethink:"binary_format,omitempty"`
	GeometryFormat interface{} `gorethink:"geometry_format,omitempty"`

	MinBatchRows              interface{} `gorethink:"min_batch_rows,omitempty"`
	MaxBatchRows              interface{} `gorethink:"max_batch_rows,omitempty"`
	MaxBatchBytes             interface{} `gorethink:"max_batch_bytes,omitempty"`
	MaxBatchSeconds           interface{} `gorethink:"max_batch_seconds,omitempty"`
	FirstBatchScaledownFactor interface{} `gorethink:"first_batch_scaledown_factor,omitempty"`
}

func (o *RunOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Run runs a query using the given connection.
//
//	rows, err := query.Run(sess)
//	if err != nil {
//		// error
//	}
//
//  var doc MyDocumentType
//	for rows.Next(&doc) {
//      // Do something with document
//	}
func (t Term) Run(s *Session, optArgs ...RunOpts) (*Cursor, error) {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}

	return s.Query(s.newQuery(t, opts))
}

// RunWrite runs a query using the given connection but unlike Run automatically
// scans the result into a variable of type WriteResponss. This function should be used
// if you are running a write query (such as Insert,  Update, TableCreate, etc...)
//
//	res, err := r.Db("database").Table("table").Insert(doc).RunWrite(sess)
func (t Term) RunWrite(s *Session, optArgs ...RunOpts) (WriteResponse, error) {
  var err error
  var res *Cursor
	var response WriteResponse

  res, err = t.Run(s, optArgs...)
	if err == nil {
		err = res.One(&response)
    defer res.Close()
	}

	return response, err
}

// ExecOpts inherits its options from RunOpts, the only difference is the
// addition of the NoReply field.
//
// When NoReply is true it causes the driver not to wait to receive the result
// and return immediately.
type ExecOpts struct {
	Db             interface{} `gorethink:"db,omitempty"`
	Profile        interface{} `gorethink:"profile,omitempty"`
	UseOutdated    interface{} `gorethink:"use_outdated,omitempty"`
	ArrayLimit     interface{} `gorethink:"array_limit,omitempty"`
	TimeFormat     interface{} `gorethink:"time_format,omitempty"`
	GroupFormat    interface{} `gorethink:"group_format,omitempty"`
	BinaryFormat   interface{} `gorethink:"binary_format,omitempty"`
	GeometryFormat interface{} `gorethink:"geometry_format,omitempty"`

	MinBatchRows              interface{} `gorethink:"min_batch_rows,omitempty"`
	MaxBatchRows              interface{} `gorethink:"max_batch_rows,omitempty"`
	MaxBatchBytes             interface{} `gorethink:"max_batch_bytes,omitempty"`
	MaxBatchSeconds           interface{} `gorethink:"max_batch_seconds,omitempty"`
	FirstBatchScaledownFactor interface{} `gorethink:"first_batch_scaledown_factor,omitempty"`

	NoReply interface{} `gorethink:"noreply,omitempty"`
}

func (o *ExecOpts) toMap() map[string]interface{} {
	return optArgsToMap(o)
}

// Exec runs the query but does not return the result. Exec will still wait for
// the response to be received unless the NoReply field is true.
//
//	res, err := r.Db("database").Table("table").Insert(doc).Exec(sess, r.ExecOpts{
//		NoReply: true,
//	})
func (t Term) Exec(s *Session, optArgs ...ExecOpts) error {
	opts := map[string]interface{}{}
	if len(optArgs) >= 1 {
		opts = optArgs[0].toMap()
	}

	return s.Exec(s.newQuery(t, opts))
}
<|endoftext|>


<commit_before>package redis // import "gopkg.in/redis.v5"

import (
	"fmt"
	"log"
	"time"

	"gopkg.in/redis.v5/internal"
	"gopkg.in/redis.v5/internal/pool"
	"gopkg.in/redis.v5/internal/proto"
)

// Redis nil reply, .e.g. when key does not exist.
const Nil = internal.Nil

func SetLogger(logger *log.Logger) {
	internal.Logger = logger
}

type baseClient struct {
	connPool pool.Pooler
	opt      *Options

	process func(Cmder) error
	onClose func() error // hook called when client is closed
}

func (c *baseClient) String() string {
	return fmt.Sprintf("Redis<%s db:%d>", c.getAddr(), c.opt.DB)
}

func (c *baseClient) conn() (*pool.Conn, bool, error) {
	cn, isNew, err := c.connPool.Get()
	if err != nil {
		return nil, false, err
	}
	if !cn.Inited {
		if err := c.initConn(cn); err != nil {
			_ = c.connPool.Remove(cn, err)
			return nil, false, err
		}
	}
	return cn, isNew, nil
}

func (c *baseClient) putConn(cn *pool.Conn, err error, allowTimeout bool) bool {
	if internal.IsBadConn(err, allowTimeout) {
		_ = c.connPool.Remove(cn, err)
		return false
	}

	_ = c.connPool.Put(cn)
	return true
}

func (c *baseClient) initConn(cn *pool.Conn) error {
	cn.Inited = true

	if c.opt.Password == "" && c.opt.DB == 0 && !c.opt.ReadOnly {
		return nil
	}

	// Temp client for Auth and Select.
	client := newClient(c.opt, pool.NewSingleConnPool(cn))
	_, err := client.Pipelined(func(pipe *Pipeline) error {
		if c.opt.Password != "" {
			pipe.Auth(c.opt.Password)
		}

		if c.opt.DB > 0 {
			pipe.Select(c.opt.DB)
		}

		if c.opt.ReadOnly {
			pipe.ReadOnly()
		}

		return nil
	})
	return err
}

func (c *baseClient) Process(cmd Cmder) error {
	if c.process != nil {
		return c.process(cmd)
	}
	return c.defaultProcess(cmd)
}

// WrapProcess replaces the process func. It takes a function createWrapper
// which is supplied by the user. createWrapper takes the old process func as
// an input and returns the new wrapper process func. createWrapper should
// use call the old process func within the new process func.
func (c *baseClient) WrapProcess(fn func(oldProcess func(cmd Cmder) error) func(cmd Cmder) error) {
	c.process = fn(c.defaultProcess)
}

func (c *baseClient) defaultProcess(cmd Cmder) error {
	for i := 0; i <= c.opt.MaxRetries; i++ {
		cn, _, err := c.conn()
		if err != nil {
			cmd.setErr(err)
			return err
		}

		cn.SetWriteTimeout(c.opt.WriteTimeout)
		if err := writeCmd(cn, cmd); err != nil {
			c.putConn(cn, err, false)
			cmd.setErr(err)
			if err != nil && internal.IsRetryableError(err) {
				continue
			}
			return err
		}

		cn.SetReadTimeout(c.cmdTimeout(cmd))
		err = cmd.readReply(cn)
		c.putConn(cn, err, false)
		if err != nil && internal.IsRetryableError(err) {
			continue
		}

		return err
	}

	return cmd.Err()
}

func (c *baseClient) cmdTimeout(cmd Cmder) time.Duration {
	if timeout := cmd.readTimeout(); timeout != nil {
		return *timeout
	} else {
		return c.opt.ReadTimeout
	}
}

func (c *baseClient) closed() bool {
	return c.connPool.Closed()
}

// Close closes the client, releasing any open resources.
//
// It is rare to Close a Client, as the Client is meant to be
// long-lived and shared between many goroutines.
func (c *baseClient) Close() error {
	var firstErr error
	if c.onClose != nil {
		if err := c.onClose(); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	if err := c.connPool.Close(); err != nil && firstErr == nil {
		firstErr = err
	}
	return firstErr
}

func (c *baseClient) getAddr() string {
	return c.opt.Addr
}

type pipelineProcessor func(*pool.Conn, []Cmder) (bool, error)

func (c *baseClient) pipelineExecer(p pipelineProcessor) pipelineExecer {
	return func(cmds []Cmder) error {
		var firstErr error
		for i := 0; i <= c.opt.MaxRetries; i++ {
			cn, _, err := c.conn()
			if err != nil {
				setCmdsErr(cmds, err)
				return err
			}

			canRetry, err := p(cn, cmds)
			c.putConn(cn, err, false)
			if err == nil {
				return nil
			}
			if firstErr == nil {
				firstErr = err
			}
			if !canRetry || !internal.IsRetryableError(err) {
				break
			}
		}
		return firstErr
	}
}

func (c *baseClient) pipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (retry bool, firstErr error) {
	cn.SetWriteTimeout(c.opt.WriteTimeout)
	if err := writeCmd(cn, cmds...); err != nil {
		setCmdsErr(cmds, err)
		return true, err
	}

	// Set read timeout for all commands.
	cn.SetReadTimeout(c.opt.ReadTimeout)
	return pipelineReadCmds(cn, cmds)
}

func pipelineReadCmds(cn *pool.Conn, cmds []Cmder) (retry bool, firstErr error) {
	for i, cmd := range cmds {
		err := cmd.readReply(cn)
		if err == nil {
			continue
		}
		if i == 0 {
			retry = true
		}
		if firstErr == nil {
			firstErr = err
		}
	}
	return false, firstErr
}

func (c *baseClient) txPipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (bool, error) {
	cn.SetWriteTimeout(c.opt.WriteTimeout)
	if err := txPipelineWriteMulti(cn, cmds); err != nil {
		setCmdsErr(cmds, err)
		return true, err
	}

	// Set read timeout for all commands.
	cn.SetReadTimeout(c.opt.ReadTimeout)

	if err := c.txPipelineReadQueued(cn, cmds); err != nil {
		return false, err
	}

	_, err := pipelineReadCmds(cn, cmds)
	return false, err
}

func txPipelineWriteMulti(cn *pool.Conn, cmds []Cmder) error {
	multiExec := make([]Cmder, 0, len(cmds)+2)
	multiExec = append(multiExec, NewStatusCmd("MULTI"))
	multiExec = append(multiExec, cmds...)
	multiExec = append(multiExec, NewSliceCmd("EXEC"))
	return writeCmd(cn, multiExec...)
}

func (c *baseClient) txPipelineReadQueued(cn *pool.Conn, cmds []Cmder) error {
	var firstErr error

	// Parse queued replies.
	var statusCmd StatusCmd
	if err := statusCmd.readReply(cn); err != nil && firstErr == nil {
		firstErr = err
	}

	for _, cmd := range cmds {
		err := statusCmd.readReply(cn)
		if err != nil {
			cmd.setErr(err)
			if firstErr == nil {
				firstErr = err
			}
		}
	}

	// Parse number of replies.
	line, err := cn.Rd.ReadLine()
	if err != nil {
		if err == Nil {
			err = TxFailedErr
		}
		return err
	}

	switch line[0] {
	case proto.ErrorReply:
		return proto.ParseErrorReply(line)
	case proto.ArrayReply:
		// ok
	default:
		err := fmt.Errorf("redis: expected '*', but got line %q", line)
		return err
	}

	return nil
}

//------------------------------------------------------------------------------

// Client is a Redis client representing a pool of zero or more
// underlying connections. It's safe for concurrent use by multiple
// goroutines.
type Client struct {
	baseClient
	cmdable
}

func newClient(opt *Options, pool pool.Pooler) *Client {
	base := baseClient{opt: opt, connPool: pool}
	client := &Client{
		baseClient: base,
		cmdable:    cmdable{base.Process},
	}
	return client
}

// NewClient returns a client to the Redis Server specified by Options.
func NewClient(opt *Options) *Client {
	opt.init()
	return newClient(opt, newConnPool(opt))
}

// PoolStats returns connection pool stats.
func (c *Client) PoolStats() *PoolStats {
	s := c.connPool.Stats()
	return &PoolStats{
		Requests: s.Requests,
		Hits:     s.Hits,
		Timeouts: s.Timeouts,

		TotalConns: s.TotalConns,
		FreeConns:  s.FreeConns,
	}
}

func (c *Client) Pipelined(fn func(*Pipeline) error) ([]Cmder, error) {
	return c.Pipeline().pipelined(fn)
}

func (c *Client) Pipeline() *Pipeline {
	pipe := Pipeline{
		exec: c.pipelineExecer(c.pipelineProcessCmds),
	}
	pipe.cmdable.process = pipe.Process
	pipe.statefulCmdable.process = pipe.Process
	return &pipe
}

func (c *Client) TxPipelined(fn func(*Pipeline) error) ([]Cmder, error) {
	return c.TxPipeline().pipelined(fn)
}

// TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
func (c *Client) TxPipeline() *Pipeline {
	pipe := Pipeline{
		exec: c.pipelineExecer(c.txPipelineProcessCmds),
	}
	pipe.cmdable.process = pipe.Process
	pipe.statefulCmdable.process = pipe.Process
	return &pipe
}

func (c *Client) pubSub() *PubSub {
	return &PubSub{
		base: baseClient{
			opt:      c.opt,
			connPool: pool.NewStickyConnPool(c.connPool.(*pool.ConnPool), false),
		},
	}
}

// Subscribe subscribes the client to the specified channels.
func (c *Client) Subscribe(channels ...string) (*PubSub, error) {
	pubsub := c.pubSub()
	return pubsub, pubsub.Subscribe(channels...)
}

// PSubscribe subscribes the client to the given patterns.
func (c *Client) PSubscribe(channels ...string) (*PubSub, error) {
	pubsub := c.pubSub()
	return pubsub, pubsub.PSubscribe(channels...)
}
<commit_msg>Allow creating PubSub without channels.<commit_after>package redis // import "gopkg.in/redis.v5"

import (
	"fmt"
	"log"
	"time"

	"gopkg.in/redis.v5/internal"
	"gopkg.in/redis.v5/internal/pool"
	"gopkg.in/redis.v5/internal/proto"
)

// Redis nil reply, .e.g. when key does not exist.
const Nil = internal.Nil

func SetLogger(logger *log.Logger) {
	internal.Logger = logger
}

type baseClient struct {
	connPool pool.Pooler
	opt      *Options

	process func(Cmder) error
	onClose func() error // hook called when client is closed
}

func (c *baseClient) String() string {
	return fmt.Sprintf("Redis<%s db:%d>", c.getAddr(), c.opt.DB)
}

func (c *baseClient) conn() (*pool.Conn, bool, error) {
	cn, isNew, err := c.connPool.Get()
	if err != nil {
		return nil, false, err
	}
	if !cn.Inited {
		if err := c.initConn(cn); err != nil {
			_ = c.connPool.Remove(cn, err)
			return nil, false, err
		}
	}
	return cn, isNew, nil
}

func (c *baseClient) putConn(cn *pool.Conn, err error, allowTimeout bool) bool {
	if internal.IsBadConn(err, allowTimeout) {
		_ = c.connPool.Remove(cn, err)
		return false
	}

	_ = c.connPool.Put(cn)
	return true
}

func (c *baseClient) initConn(cn *pool.Conn) error {
	cn.Inited = true

	if c.opt.Password == "" && c.opt.DB == 0 && !c.opt.ReadOnly {
		return nil
	}

	// Temp client for Auth and Select.
	client := newClient(c.opt, pool.NewSingleConnPool(cn))
	_, err := client.Pipelined(func(pipe *Pipeline) error {
		if c.opt.Password != "" {
			pipe.Auth(c.opt.Password)
		}

		if c.opt.DB > 0 {
			pipe.Select(c.opt.DB)
		}

		if c.opt.ReadOnly {
			pipe.ReadOnly()
		}

		return nil
	})
	return err
}

func (c *baseClient) Process(cmd Cmder) error {
	if c.process != nil {
		return c.process(cmd)
	}
	return c.defaultProcess(cmd)
}

// WrapProcess replaces the process func. It takes a function createWrapper
// which is supplied by the user. createWrapper takes the old process func as
// an input and returns the new wrapper process func. createWrapper should
// use call the old process func within the new process func.
func (c *baseClient) WrapProcess(fn func(oldProcess func(cmd Cmder) error) func(cmd Cmder) error) {
	c.process = fn(c.defaultProcess)
}

func (c *baseClient) defaultProcess(cmd Cmder) error {
	for i := 0; i <= c.opt.MaxRetries; i++ {
		cn, _, err := c.conn()
		if err != nil {
			cmd.setErr(err)
			return err
		}

		cn.SetWriteTimeout(c.opt.WriteTimeout)
		if err := writeCmd(cn, cmd); err != nil {
			c.putConn(cn, err, false)
			cmd.setErr(err)
			if err != nil && internal.IsRetryableError(err) {
				continue
			}
			return err
		}

		cn.SetReadTimeout(c.cmdTimeout(cmd))
		err = cmd.readReply(cn)
		c.putConn(cn, err, false)
		if err != nil && internal.IsRetryableError(err) {
			continue
		}

		return err
	}

	return cmd.Err()
}

func (c *baseClient) cmdTimeout(cmd Cmder) time.Duration {
	if timeout := cmd.readTimeout(); timeout != nil {
		return *timeout
	} else {
		return c.opt.ReadTimeout
	}
}

func (c *baseClient) closed() bool {
	return c.connPool.Closed()
}

// Close closes the client, releasing any open resources.
//
// It is rare to Close a Client, as the Client is meant to be
// long-lived and shared between many goroutines.
func (c *baseClient) Close() error {
	var firstErr error
	if c.onClose != nil {
		if err := c.onClose(); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	if err := c.connPool.Close(); err != nil && firstErr == nil {
		firstErr = err
	}
	return firstErr
}

func (c *baseClient) getAddr() string {
	return c.opt.Addr
}

type pipelineProcessor func(*pool.Conn, []Cmder) (bool, error)

func (c *baseClient) pipelineExecer(p pipelineProcessor) pipelineExecer {
	return func(cmds []Cmder) error {
		var firstErr error
		for i := 0; i <= c.opt.MaxRetries; i++ {
			cn, _, err := c.conn()
			if err != nil {
				setCmdsErr(cmds, err)
				return err
			}

			canRetry, err := p(cn, cmds)
			c.putConn(cn, err, false)
			if err == nil {
				return nil
			}
			if firstErr == nil {
				firstErr = err
			}
			if !canRetry || !internal.IsRetryableError(err) {
				break
			}
		}
		return firstErr
	}
}

func (c *baseClient) pipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (retry bool, firstErr error) {
	cn.SetWriteTimeout(c.opt.WriteTimeout)
	if err := writeCmd(cn, cmds...); err != nil {
		setCmdsErr(cmds, err)
		return true, err
	}

	// Set read timeout for all commands.
	cn.SetReadTimeout(c.opt.ReadTimeout)
	return pipelineReadCmds(cn, cmds)
}

func pipelineReadCmds(cn *pool.Conn, cmds []Cmder) (retry bool, firstErr error) {
	for i, cmd := range cmds {
		err := cmd.readReply(cn)
		if err == nil {
			continue
		}
		if i == 0 {
			retry = true
		}
		if firstErr == nil {
			firstErr = err
		}
	}
	return false, firstErr
}

func (c *baseClient) txPipelineProcessCmds(cn *pool.Conn, cmds []Cmder) (bool, error) {
	cn.SetWriteTimeout(c.opt.WriteTimeout)
	if err := txPipelineWriteMulti(cn, cmds); err != nil {
		setCmdsErr(cmds, err)
		return true, err
	}

	// Set read timeout for all commands.
	cn.SetReadTimeout(c.opt.ReadTimeout)

	if err := c.txPipelineReadQueued(cn, cmds); err != nil {
		return false, err
	}

	_, err := pipelineReadCmds(cn, cmds)
	return false, err
}

func txPipelineWriteMulti(cn *pool.Conn, cmds []Cmder) error {
	multiExec := make([]Cmder, 0, len(cmds)+2)
	multiExec = append(multiExec, NewStatusCmd("MULTI"))
	multiExec = append(multiExec, cmds...)
	multiExec = append(multiExec, NewSliceCmd("EXEC"))
	return writeCmd(cn, multiExec...)
}

func (c *baseClient) txPipelineReadQueued(cn *pool.Conn, cmds []Cmder) error {
	var firstErr error

	// Parse queued replies.
	var statusCmd StatusCmd
	if err := statusCmd.readReply(cn); err != nil && firstErr == nil {
		firstErr = err
	}

	for _, cmd := range cmds {
		err := statusCmd.readReply(cn)
		if err != nil {
			cmd.setErr(err)
			if firstErr == nil {
				firstErr = err
			}
		}
	}

	// Parse number of replies.
	line, err := cn.Rd.ReadLine()
	if err != nil {
		if err == Nil {
			err = TxFailedErr
		}
		return err
	}

	switch line[0] {
	case proto.ErrorReply:
		return proto.ParseErrorReply(line)
	case proto.ArrayReply:
		// ok
	default:
		err := fmt.Errorf("redis: expected '*', but got line %q", line)
		return err
	}

	return nil
}

//------------------------------------------------------------------------------

// Client is a Redis client representing a pool of zero or more
// underlying connections. It's safe for concurrent use by multiple
// goroutines.
type Client struct {
	baseClient
	cmdable
}

func newClient(opt *Options, pool pool.Pooler) *Client {
	base := baseClient{opt: opt, connPool: pool}
	client := &Client{
		baseClient: base,
		cmdable:    cmdable{base.Process},
	}
	return client
}

// NewClient returns a client to the Redis Server specified by Options.
func NewClient(opt *Options) *Client {
	opt.init()
	return newClient(opt, newConnPool(opt))
}

// PoolStats returns connection pool stats.
func (c *Client) PoolStats() *PoolStats {
	s := c.connPool.Stats()
	return &PoolStats{
		Requests: s.Requests,
		Hits:     s.Hits,
		Timeouts: s.Timeouts,

		TotalConns: s.TotalConns,
		FreeConns:  s.FreeConns,
	}
}

func (c *Client) Pipelined(fn func(*Pipeline) error) ([]Cmder, error) {
	return c.Pipeline().pipelined(fn)
}

func (c *Client) Pipeline() *Pipeline {
	pipe := Pipeline{
		exec: c.pipelineExecer(c.pipelineProcessCmds),
	}
	pipe.cmdable.process = pipe.Process
	pipe.statefulCmdable.process = pipe.Process
	return &pipe
}

func (c *Client) TxPipelined(fn func(*Pipeline) error) ([]Cmder, error) {
	return c.TxPipeline().pipelined(fn)
}

// TxPipeline acts like Pipeline, but wraps queued commands with MULTI/EXEC.
func (c *Client) TxPipeline() *Pipeline {
	pipe := Pipeline{
		exec: c.pipelineExecer(c.txPipelineProcessCmds),
	}
	pipe.cmdable.process = pipe.Process
	pipe.statefulCmdable.process = pipe.Process
	return &pipe
}

func (c *Client) pubSub() *PubSub {
	return &PubSub{
		base: baseClient{
			opt:      c.opt,
			connPool: pool.NewStickyConnPool(c.connPool.(*pool.ConnPool), false),
		},
	}
}

// Subscribe subscribes the client to the specified channels.
func (c *Client) Subscribe(channels ...string) (*PubSub, error) {
	pubsub := c.pubSub()
	if len(channels) > 0 {
		if err := pubsub.Subscribe(channels...); err != nil {
			pubsub.Close()
			return nil, err
		}
	}
	return pubsub, nil
}

// PSubscribe subscribes the client to the given patterns.
func (c *Client) PSubscribe(channels ...string) (*PubSub, error) {
	pubsub := c.pubSub()
	if len(channels) > 0 {
		if err := pubsub.PSubscribe(channels...); err != nil {
			pubsub.Close()
			return nil, err
		}
	}
	return pubsub, nil
}
<|endoftext|>