<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"encoding/json"
	"os"
	"os/signal"
	"regexp"
	"runtime"
	"syscall"
	"strconv"
	"time"
	"io"
	"net/http"
	"github.com/hoisie/redis"
	"github.com/gorilla/mux"
)

const (
	VERSION                 = "0.1.0"
	MAX_UNPROCESSED_PACKETS = 1000
	MAX_LOG_ENTRIES         = 1000
	MAX_UDP_PACKET_SIZE     = 512
)

var (
	serviceAddress   = flag.String("address", ":8125", "UDP service address")
	expiryInterval   = flag.Int64("expiry-interval", 1, "Expiry interval (seconds)")
	debug            = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion      = flag.Bool("version", false, "print version string")
)

const (
	ACTION_SET_WARN = "set-warn"
	ACTION_SET_ERR = "set-err"
	ACTION_BEAT = "beat"
)

const (
	STATE_PAUSED  = "paused"
	STATE_OK      = "ok"
	STATE_WARNING = "warning"
	STATE_ERROR   = "error"
)

type Cmd struct {
	Action   string
	Service  string
	Value    int
}

type Service struct {
	Name           string
	LastValue      int
	LastBeat       int64
	LastUpdated    int64
	WarningTimeout int64
	ErrorTimeout   int64
	State          string
}

var (
	In       = make(chan *Cmd, MAX_UNPROCESSED_PACKETS)
	signalchan chan os.Signal
	client redis.Client
)

func now() int64 { return time.Now().Unix() }

func getOrCreate(name string) (*Service, *Service) {
	service := &Service{
		Name: name,
		LastValue: -1,
		LastBeat: -1,
		LastUpdated: -1,
		WarningTimeout: -1,
		ErrorTimeout: -1,
		State: STATE_PAUSED,
	}

	if data, err := client.Get("lb.service." + name); err == nil {
		json.Unmarshal(data, &service)
	}
	var ref = *service
	return service, &ref
}

func (s *Service)GetExpiry(timeout int64) int64 {
	if timeout <= 0 {
		return 0
	}
	return s.LastBeat + timeout
}

func (s *Service)GetNextExpiry(ts int64) int64 {
	var next int64 = 0
	var warningExpiry = s.GetExpiry(s.WarningTimeout)
	var errorExpiry = s.GetExpiry(s.ErrorTimeout)
	if warningExpiry > 0 && warningExpiry > ts && (next == 0 || warningExpiry < next) {
		next = warningExpiry
	}
	if errorExpiry > 0 && errorExpiry > ts && (next == 0 || errorExpiry < next) {
		next = errorExpiry
	}
	return next
}

func (s *Service) UpdateState(ts int64) {
	s.State = STATE_OK
	var warningExpiry = s.GetExpiry(s.WarningTimeout)
	var errorExpiry = s.GetExpiry(s.ErrorTimeout)
	if warningExpiry > 0 && ts >= warningExpiry {
		s.State = STATE_WARNING
	}
	if errorExpiry > 0 && ts >= errorExpiry {
		s.State = STATE_ERROR
	}
}

func (s *Service) Log(format string, args ...interface{}) {
	var key = "lb.service-log." + s.Name
	var log = fmt.Sprintf(format, args...)
	client.Lpush(key, []byte(log))
	client.Ltrim(key, 0, MAX_LOG_ENTRIES)
}

func updateExpiry(service *Service, ts int64) {
	if service.State != STATE_PAUSED {
		if expiry := service.GetNextExpiry(ts); expiry > 0 {
			client.Zadd("lb.expiry", []byte(service.Name), float64(expiry))
			return
		}
	}
	client.Zrem("lb.expiry", []byte(service.Name))
}

func (s *Service)Save(ref *Service, ts int64) {
	if *s != *ref {
		if s.State != ref.State {
			log.Printf("service %s, state %s -> %s",
				s.Name, ref.State, s.State)
			s.Log("%d|state|%s", ts, s.State)
		}
		if s.WarningTimeout != ref.WarningTimeout {
			log.Printf("service %s, warn %d -> %d",
				s.Name, ref.WarningTimeout, s.WarningTimeout)
			s.Log("%d|warn-tmo|%s", ts, ref.WarningTimeout)
		}
		if s.ErrorTimeout != ref.ErrorTimeout {
			log.Printf("service %s, err %d -> %d",
				s.Name, ref.ErrorTimeout, s.ErrorTimeout)
			s.Log("%d|err-tmo|%s", ts, ref.ErrorTimeout)
		}
		s.LastUpdated = ts
		b, _ := json.Marshal(s)
		client.Set("lb.service." + s.Name, b)
		if ref.LastUpdated < 0 {
			client.Sadd("lb.services.all", []byte(s.Name))
		}
	}
}

func monitor() {
	period := time.Duration(*expiryInterval) * time.Second
	ticker := time.NewTicker(period)
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			return
		case <-ticker.C:
			var ts = now()
			if expired, err := client.Zrangebyscore("lb.expiry", 0, float64(now())); err == nil {
				for _, elem := range expired {
					var service, ref = getOrCreate(string(elem))
					service.UpdateState(ts)
					service.Save(ref, ts)
					updateExpiry(service, ts)
				}
			}
		case s := <-In:
			var ts = now()
			var service, ref = getOrCreate(s.Service)
			switch s.Action {
			case ACTION_SET_WARN:
				service.WarningTimeout = int64(s.Value)
			case ACTION_SET_ERR:
				service.ErrorTimeout = int64(s.Value)
			case ACTION_BEAT:
				service.LastBeat = ts
				var diff = ts - ref.LastBeat
				service.Log("%d|beat|%d", ts, diff)
			}
			service.UpdateState(ts)
			service.Save(ref, ts)
			updateExpiry(service, ts)
		}
	}
}

var packetRegexp = regexp.MustCompile("^([^:]+)\\.(beat|warn|err):(-?[0-9]+)\\|(g|c|ms)(\\|@([0-9\\.]+))?\n?$")

func parseMessage(data []byte) []*Cmd {
	var output []*Cmd
	for _, line := range bytes.Split(data, []byte("\n")) {
		if len(line) == 0 {
			continue
		}

		item := packetRegexp.FindSubmatch(line)
		if len(item) == 0 {
			continue
		}

		var value int
		modifier := string(item[4])
		switch modifier {
		case "c":
			var vali, err = strconv.ParseInt(string(item[3]), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseInt %s - %s", item[3], err)
				continue
			}
			value = int(vali)
		default:
			var valu, err = strconv.ParseUint(string(item[3]), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseUint %s - %s", item[3], err)
				continue
			}
			value = int(valu)
		}
		var action string
		switch string(item[2]) {
		case "warn":
			action = ACTION_SET_WARN
		case "err":
			action = ACTION_SET_ERR
		case "beat":
			action = ACTION_BEAT
		}

		packet := &Cmd{
			Action: action,
			Service: string(item[1]),
			Value:    value,
		}
		output = append(output, packet)
	}
	return output
}

func udpListener() {
	address, _ := net.ResolveUDPAddr("udp", *serviceAddress)
	log.Printf("listening on %s", address)
	listener, err := net.ListenUDP("udp", address)
	if err != nil {
		log.Fatalf("ERROR: ListenUDP - %s", err)
	}
	defer listener.Close()

	message := make([]byte, MAX_UDP_PACKET_SIZE)
	for {
		n, remaddr, err := listener.ReadFromUDP(message)
		if err != nil {
			log.Printf("ERROR: reading UDP packet from %+v - %s", remaddr, err)
			continue
		}

		for _, p := range parseMessage(message[:n]) {
			In <- p
		}
	}
}

func DashboardHandler(c http.ResponseWriter, req *http.Request) {
	var buffer bytes.Buffer
	var services, _ = client.Smembers("lb.services.all")
	var errors, warnings = false, false
	for _, v := range services {
		var service, _ = getOrCreate(string(v))
		buffer.WriteString(fmt.Sprintf("%-40s%s\n", service.Name, service.State))
		if service.State == STATE_WARNING {
			warnings = true
		}
		if service.State == STATE_ERROR {
			errors = true
		}
	}
	buffer.WriteString(fmt.Sprintf("\nwarnings: %t\nerrors: %t\ngood: %t\n", warnings, errors, !warnings && !errors))
        body := buffer.String()
        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", strconv.Itoa(len(body)))
        io.WriteString(c, body)
}

func StatusHandler(c http.ResponseWriter, req *http.Request) {
	var buffer bytes.Buffer
	var services, _ = client.Smembers("lb.services.all")
	var errors, warnings = false, false
	for _, v := range services {
		var service, _ = getOrCreate(string(v))
		if service.State == STATE_WARNING {
			warnings = true
		}
		if service.State == STATE_ERROR {
			errors = true
		}
	}
	buffer.WriteString(fmt.Sprintf("warnings: %t\nerrors: %t\ngood: %t\n", warnings, errors, !warnings && !errors))
        body := buffer.String()
        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", strconv.Itoa(len(body)))
        io.WriteString(c, body)
}

func TriggerHandler(c http.ResponseWriter, r *http.Request) {
	params := mux.Vars(r)
	name := params["name"]

	var err = r.ParseForm()
	if err != nil {
		log.Print("error parsing form ", err)
		return
	}

	var errtmo, warntmo = r.FormValue("err-tmo"), r.FormValue("warn-tmo")

	In <- &Cmd{
		Action:  ACTION_BEAT,
		Service: name,
		Value:   1,
	}

	if val, err := strconv.Atoi(errtmo); err == nil {
		In <- &Cmd{
			Action:  ACTION_SET_ERR,
			Service: name,
			Value:   val,
		}
	}

	if val, err := strconv.Atoi(warntmo); err == nil {
		In <- &Cmd{
			Action:  ACTION_SET_WARN,
			Service: name,
			Value:   val,
		}
	}

        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", "3")
        io.WriteString(c, "ok\n")
}

func httpServer() {
	rtr := mux.NewRouter()
	rtr.HandleFunc("/", DashboardHandler).Methods("GET")
	rtr.HandleFunc("/status", StatusHandler).Methods("GET")
	rtr.HandleFunc("/trigger/{name:[a-z.]+}", TriggerHandler).Methods("POST")
	http.Handle("/", rtr)
        http.ListenAndServe(":8080", nil)
}

func main() {
	flag.Parse()
	if *showVersion {
		fmt.Printf("statsdaemon v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}

	signalchan = make(chan os.Signal, 1)
	signal.Notify(signalchan, syscall.SIGTERM)

	go httpServer()
	go udpListener()
	monitor()
}
<commit_msg>Different formatting of dashboard<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"encoding/json"
	"os"
	"os/signal"
	"regexp"
	"runtime"
	"syscall"
	"strconv"
	"time"
	"io"
	"net/http"
	"github.com/hoisie/redis"
	"github.com/gorilla/mux"
)

const (
	VERSION                 = "0.1.0"
	MAX_UNPROCESSED_PACKETS = 1000
	MAX_LOG_ENTRIES         = 1000
	MAX_UDP_PACKET_SIZE     = 512
)

var (
	serviceAddress   = flag.String("address", ":8125", "UDP service address")
	expiryInterval   = flag.Int64("expiry-interval", 1, "Expiry interval (seconds)")
	debug            = flag.Bool("debug", false, "print statistics sent to graphite")
	showVersion      = flag.Bool("version", false, "print version string")
)

const (
	ACTION_SET_WARN = "set-warn"
	ACTION_SET_ERR = "set-err"
	ACTION_BEAT = "beat"
)

const (
	STATE_PAUSED  = "paused"
	STATE_OK      = "ok"
	STATE_WARNING = "warning"
	STATE_ERROR   = "error"
)

type Cmd struct {
	Action   string
	Service  string
	Value    int
}

type Service struct {
	Name           string
	LastValue      int
	LastBeat       int64
	LastUpdated    int64
	WarningTimeout int64
	ErrorTimeout   int64
	State          string
}

var (
	In       = make(chan *Cmd, MAX_UNPROCESSED_PACKETS)
	signalchan chan os.Signal
	client redis.Client
)

func now() int64 { return time.Now().Unix() }

func getOrCreate(name string) (*Service, *Service) {
	service := &Service{
		Name: name,
		LastValue: -1,
		LastBeat: -1,
		LastUpdated: -1,
		WarningTimeout: -1,
		ErrorTimeout: -1,
		State: STATE_PAUSED,
	}

	if data, err := client.Get("lb.service." + name); err == nil {
		json.Unmarshal(data, &service)
	}
	var ref = *service
	return service, &ref
}

func (s *Service)GetExpiry(timeout int64) int64 {
	if timeout <= 0 {
		return 0
	}
	return s.LastBeat + timeout
}

func (s *Service)GetNextExpiry(ts int64) int64 {
	var next int64 = 0
	var warningExpiry = s.GetExpiry(s.WarningTimeout)
	var errorExpiry = s.GetExpiry(s.ErrorTimeout)
	if warningExpiry > 0 && warningExpiry > ts && (next == 0 || warningExpiry < next) {
		next = warningExpiry
	}
	if errorExpiry > 0 && errorExpiry > ts && (next == 0 || errorExpiry < next) {
		next = errorExpiry
	}
	return next
}

func (s *Service) UpdateState(ts int64) {
	s.State = STATE_OK
	var warningExpiry = s.GetExpiry(s.WarningTimeout)
	var errorExpiry = s.GetExpiry(s.ErrorTimeout)
	if warningExpiry > 0 && ts >= warningExpiry {
		s.State = STATE_WARNING
	}
	if errorExpiry > 0 && ts >= errorExpiry {
		s.State = STATE_ERROR
	}
}

func (s *Service) Log(format string, args ...interface{}) {
	var key = "lb.service-log." + s.Name
	var log = fmt.Sprintf(format, args...)
	client.Lpush(key, []byte(log))
	client.Ltrim(key, 0, MAX_LOG_ENTRIES)
}

func updateExpiry(service *Service, ts int64) {
	if service.State != STATE_PAUSED {
		if expiry := service.GetNextExpiry(ts); expiry > 0 {
			client.Zadd("lb.expiry", []byte(service.Name), float64(expiry))
			return
		}
	}
	client.Zrem("lb.expiry", []byte(service.Name))
}

func (s *Service)Save(ref *Service, ts int64) {
	if *s != *ref {
		if s.State != ref.State {
			log.Printf("service %s, state %s -> %s",
				s.Name, ref.State, s.State)
			s.Log("%d|state|%s", ts, s.State)
		}
		if s.WarningTimeout != ref.WarningTimeout {
			log.Printf("service %s, warn %d -> %d",
				s.Name, ref.WarningTimeout, s.WarningTimeout)
			s.Log("%d|warn-tmo|%s", ts, ref.WarningTimeout)
		}
		if s.ErrorTimeout != ref.ErrorTimeout {
			log.Printf("service %s, err %d -> %d",
				s.Name, ref.ErrorTimeout, s.ErrorTimeout)
			s.Log("%d|err-tmo|%s", ts, ref.ErrorTimeout)
		}
		s.LastUpdated = ts
		b, _ := json.Marshal(s)
		client.Set("lb.service." + s.Name, b)
		if ref.LastUpdated < 0 {
			client.Sadd("lb.services.all", []byte(s.Name))
		}
	}
}

func monitor() {
	period := time.Duration(*expiryInterval) * time.Second
	ticker := time.NewTicker(period)
	for {
		select {
		case sig := <-signalchan:
			fmt.Printf("!! Caught signal %d... shutting down\n", sig)
			return
		case <-ticker.C:
			var ts = now()
			if expired, err := client.Zrangebyscore("lb.expiry", 0, float64(now())); err == nil {
				for _, elem := range expired {
					var service, ref = getOrCreate(string(elem))
					service.UpdateState(ts)
					service.Save(ref, ts)
					updateExpiry(service, ts)
				}
			}
		case s := <-In:
			var ts = now()
			var service, ref = getOrCreate(s.Service)
			switch s.Action {
			case ACTION_SET_WARN:
				service.WarningTimeout = int64(s.Value)
			case ACTION_SET_ERR:
				service.ErrorTimeout = int64(s.Value)
			case ACTION_BEAT:
				service.LastBeat = ts
				var diff = ts - ref.LastBeat
				service.Log("%d|beat|%d", ts, diff)
			}
			service.UpdateState(ts)
			service.Save(ref, ts)
			updateExpiry(service, ts)
		}
	}
}

var packetRegexp = regexp.MustCompile("^([^:]+)\\.(beat|warn|err):(-?[0-9]+)\\|(g|c|ms)(\\|@([0-9\\.]+))?\n?$")

func parseMessage(data []byte) []*Cmd {
	var output []*Cmd
	for _, line := range bytes.Split(data, []byte("\n")) {
		if len(line) == 0 {
			continue
		}

		item := packetRegexp.FindSubmatch(line)
		if len(item) == 0 {
			continue
		}

		var value int
		modifier := string(item[4])
		switch modifier {
		case "c":
			var vali, err = strconv.ParseInt(string(item[3]), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseInt %s - %s", item[3], err)
				continue
			}
			value = int(vali)
		default:
			var valu, err = strconv.ParseUint(string(item[3]), 10, 64)
			if err != nil {
				log.Printf("ERROR: failed to ParseUint %s - %s", item[3], err)
				continue
			}
			value = int(valu)
		}
		var action string
		switch string(item[2]) {
		case "warn":
			action = ACTION_SET_WARN
		case "err":
			action = ACTION_SET_ERR
		case "beat":
			action = ACTION_BEAT
		}

		packet := &Cmd{
			Action: action,
			Service: string(item[1]),
			Value:    value,
		}
		output = append(output, packet)
	}
	return output
}

func udpListener() {
	address, _ := net.ResolveUDPAddr("udp", *serviceAddress)
	log.Printf("listening on %s", address)
	listener, err := net.ListenUDP("udp", address)
	if err != nil {
		log.Fatalf("ERROR: ListenUDP - %s", err)
	}
	defer listener.Close()

	message := make([]byte, MAX_UDP_PACKET_SIZE)
	for {
		n, remaddr, err := listener.ReadFromUDP(message)
		if err != nil {
			log.Printf("ERROR: reading UDP packet from %+v - %s", remaddr, err)
			continue
		}

		for _, p := range parseMessage(message[:n]) {
			In <- p
		}
	}
}

func DashboardState(in string) string {
	return map[string]string {
		STATE_WARNING: "  WARN   ",
		STATE_ERROR:   "      ERR",
		STATE_OK:      "OK       ",
	}[in]
}

func DashboardHandler(c http.ResponseWriter, req *http.Request) {
	var buffer bytes.Buffer
	var services, _ = client.Smembers("lb.services.all")
	var errors, warnings = false, false

	for _, v := range services {
		var service, _ = getOrCreate(string(v))
		buffer.WriteString(fmt.Sprintf(" %s | %-40s\n", DashboardState(service.State), service.Name))
		if service.State == STATE_WARNING {
			warnings = true
		}
		if service.State == STATE_ERROR {
			errors = true
		}
	}
	buffer.WriteString(fmt.Sprintf("\nwarnings: %t\nerrors: %t\ngood: %t\n", warnings, errors, !warnings && !errors))
        body := buffer.String()
        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", strconv.Itoa(len(body)))
        io.WriteString(c, body)
}

func StatusHandler(c http.ResponseWriter, req *http.Request) {
	var buffer bytes.Buffer
	var services, _ = client.Smembers("lb.services.all")
	var errors, warnings = false, false
	for _, v := range services {
		var service, _ = getOrCreate(string(v))
		if service.State == STATE_WARNING {
			warnings = true
		}
		if service.State == STATE_ERROR {
			errors = true
		}
	}
	buffer.WriteString(fmt.Sprintf("warnings: %t\nerrors: %t\ngood: %t\n", warnings, errors, !warnings && !errors))
        body := buffer.String()
        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", strconv.Itoa(len(body)))
        io.WriteString(c, body)
}

func TriggerHandler(c http.ResponseWriter, r *http.Request) {
	params := mux.Vars(r)
	name := params["name"]

	var err = r.ParseForm()
	if err != nil {
		log.Print("error parsing form ", err)
		return
	}

	var errtmo, warntmo = r.FormValue("err-tmo"), r.FormValue("warn-tmo")

	In <- &Cmd{
		Action:  ACTION_BEAT,
		Service: name,
		Value:   1,
	}

	if val, err := strconv.Atoi(errtmo); err == nil {
		In <- &Cmd{
			Action:  ACTION_SET_ERR,
			Service: name,
			Value:   val,
		}
	}

	if val, err := strconv.Atoi(warntmo); err == nil {
		In <- &Cmd{
			Action:  ACTION_SET_WARN,
			Service: name,
			Value:   val,
		}
	}

        c.Header().Add("Content-Type", "text/plain")
        c.Header().Add("Content-Length", "3")
        io.WriteString(c, "ok\n")
}

func httpServer() {
	rtr := mux.NewRouter()
	rtr.HandleFunc("/", DashboardHandler).Methods("GET")
	rtr.HandleFunc("/status", StatusHandler).Methods("GET")
	rtr.HandleFunc("/trigger/{name:[a-z.]+}", TriggerHandler).Methods("POST")
	http.Handle("/", rtr)
        http.ListenAndServe(":8080", nil)
}

func main() {
	flag.Parse()
	if *showVersion {
		fmt.Printf("statsdaemon v%s (built w/%s)\n", VERSION, runtime.Version())
		return
	}

	signalchan = make(chan os.Signal, 1)
	signal.Notify(signalchan, syscall.SIGTERM)

	go httpServer()
	go udpListener()
	monitor()
}
<|endoftext|>


<commit_before>// Copyright (c) 2003-2005 Maxim Sobolev. All rights reserved.
// Copyright (c) 2006-2014 Sippy Software, Inc. All rights reserved.
// Copyright (c) 2016 Andriy Pylypenko. All rights reserved.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package sippy

import (
    "fmt"
    "os"
    "net"
    "strings"
    "syscall"
    "time"

    "sippy/conf"
    "sippy/math"
    "sippy/time"
)

const (
    _RTPPLWorker_MAX_RETRIES = 3
)

type rtpp_req_stream struct {
    command         string
    result_callback func(string)
}

type _RTPPLWorker struct {
    userv           *Rtp_proxy_client_stream
    s               net.Conn
    shutdown_chan   chan int
}

func NewRTPPLWorker(userv *Rtp_proxy_client_stream) *_RTPPLWorker {
    self := &_RTPPLWorker{
        userv           : userv,
        shutdown_chan   : make(chan int, 1),
    }
    self.reconnect()
    go self.run()
    return self
}

func (self *_RTPPLWorker) reconnect() (err error) {
    if self.s != nil {
        self.s.Close()
    }
    self.s, err = net.Dial(self.userv.address.Network(), self.userv.address.String())
    return
}

func (self *_RTPPLWorker) send_raw(command string, stime *sippy_time.MonoTime) (string, time.Duration, error) {
    if command[len(command)-1] != '\n' {
        command += "\n"
    }
    //print "%s.send_raw(%s)" % (id(self), command)
    if stime == nil {
        stime, _ = sippy_time.NewMonoTime()
    }
    var err error
    var n int
    retries := 0
    rval := ""
    buf := make([]byte, 1024)
    for {
        if retries > _RTPPLWorker_MAX_RETRIES {
            return "", 0, fmt.Errorf("Error sending to the rtpproxy on " + self.userv.address.String() + ": " + err.Error())
        }
        retries++

        if self.s == nil || err != nil {
            err = self.reconnect()
            if err != nil {
                time.Sleep(100 * time.Millisecond)
                continue
            }
        }
        _, err = self.s.Write([]byte(command))
        if err != nil {
            time.Sleep(100 * time.Millisecond)
            continue
        }
        n, err = self.s.Read(buf)
        if err != nil {
            time.Sleep(100 * time.Millisecond)
            continue
        }
        rval = strings.TrimSpace(string(buf[:n]))
        break
    }
    rtpc_delay, _ := stime.OffsetFromNow()
    return rval, rtpc_delay, nil
}

func (self *_RTPPLWorker) run() {
    for {
        req := <-self.userv.wi
        if req == nil {
            // Shutdown request, relay it further
            self.userv.wi <- nil
            break
        }
        //command, result_callback, callback_parameters = wi
        data, rtpc_delay, err := self.send_raw(req.command, nil)
        if err != nil {
            self.userv.owner.logger.Debug("Error communicating the rtpproxy: " + err.Error())
            data, rtpc_delay = "", -1
        }
        if len(data) == 0 {
            rtpc_delay = -1
        }
        if req.result_callback != nil {
            go req.result_callback(data)
        }
        if rtpc_delay != -1 {
            go self.userv.register_delay(rtpc_delay)
        }
    }
    self.shutdown_chan <- 1
}

type Rtp_proxy_client_stream struct {
    owner       *Rtp_proxy_client_base
    address     net.Addr
    nworkers    int
    workers     []*_RTPPLWorker
    delay_flt   sippy_math.RecFilter
    _is_local    bool
    wi          chan *rtpp_req_stream
}

func NewRtp_proxy_client_stream(owner *Rtp_proxy_client_base, global_config sippy_conf.Config, address net.Addr, opts *Rtp_proxy_opts) (rtp_proxy_transport, error) {
    var err error
    if address == nil {
        address, err = net.ResolveUnixAddr("unix", "/var/run/rtpproxy.sock")
        if err != nil {
            return nil, err
        }
    }
    nworkers := 4
    if opts != nil && opts.Nworkers != nil {
        nworkers = *opts.Nworkers
    }
    self := &Rtp_proxy_client_stream{
        owner       : owner,
        address     : address,
        nworkers    : nworkers,
        workers     : make([]*_RTPPLWorker, nworkers),
        delay_flt   : sippy_math.NewRecFilter(0.95, 0.25),
        wi          : make(chan *rtpp_req_stream, 1000),
    }
    if strings.HasPrefix(address.Network(), "unix") {
        self._is_local = true
    } else {
        self._is_local = false
    }
    //self.wi_available = Condition()
    //self.wi = []
    for i := 0; i < self.nworkers; i++ {
        self.workers[i] = NewRTPPLWorker(self)
    }
    return self, nil
}

func (self *Rtp_proxy_client_stream) is_local() bool {
    return self._is_local
}

func (self *Rtp_proxy_client_stream) send_command(command string, result_callback func(string)) {
    if command[len(command)-1] != '\n' {
        command += "\n"
    }
    self.wi <- &rtpp_req_stream{command, result_callback }
}
/*
    def reconnect(self, address, bind_address = nil):
        self.shutdown()
        self.address = address
        self.workers = []
        for i in range(0, self.nworkers):
            self.workers.append(_RTPPLWorker(self))
        self.delay_flt = recfilter(0.95, 0.25)
*/
func (self *Rtp_proxy_client_stream) shutdown() {
    self.wi <- nil
    for _, rworker := range self.workers {
        <-rworker.shutdown_chan
    }
    self.workers = nil
}

func (self *Rtp_proxy_client_stream) register_delay(rtpc_delay time.Duration) {
    self.delay_flt.Apply(rtpc_delay.Seconds())
}
/*
    def get_rtpc_delay(self):
        return self.delay_flt.lastval

if __name__ == "__main__":
    from twisted.internet import reactor
    def display(*args):
        print args
        reactor.crash()
    r = Rtp_proxy_client_stream({"_sip_address":"1.2.3.4"})
    r.send_command("VF 123456", display, "abcd")
    reactor.run(installSignalHandlers = 1)
    r.shutdown()
*/

func _net_errno(err error) syscall.Errno {
    oe, ok := err.(*net.OpError)
    if ! ok { return 0 }
    errno, ok := oe.Err.(*os.SyscallError)
    if ! ok { return 0 }
    ret, ok := errno.Err.(syscall.Errno)
    if ! ok { return 0 }
    return ret
}
<commit_msg>The rtpproxy closes the unix command socket after every command. Do not try to keep the opened connection.<commit_after>// Copyright (c) 2003-2005 Maxim Sobolev. All rights reserved.
// Copyright (c) 2006-2014 Sippy Software, Inc. All rights reserved.
// Copyright (c) 2016 Andriy Pylypenko. All rights reserved.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package sippy

import (
    "fmt"
    "os"
    "net"
    "strings"
    "syscall"
    "time"

    "sippy/conf"
    "sippy/math"
    "sippy/time"
)

const (
    _RTPPLWorker_MAX_RETRIES = 3
)

type rtpp_req_stream struct {
    command         string
    result_callback func(string)
}

type _RTPPLWorker struct {
    userv           *Rtp_proxy_client_stream
    shutdown_chan   chan int
}

func NewRTPPLWorker(userv *Rtp_proxy_client_stream) *_RTPPLWorker {
    self := &_RTPPLWorker{
        userv           : userv,
        shutdown_chan   : make(chan int, 1),
    }
    go self.run()
    return self
}

func (self *_RTPPLWorker) send_raw(command string, stime *sippy_time.MonoTime) (string, time.Duration, error) {
    if command[len(command)-1] != '\n' {
        command += "\n"
    }
    //print "%s.send_raw(%s)" % (id(self), command)
    if stime == nil {
        stime, _ = sippy_time.NewMonoTime()
    }
    var err error
    var n int
    retries := 0
    rval := ""
    buf := make([]byte, 1024)
    var s net.Conn
    for {
        if retries > _RTPPLWorker_MAX_RETRIES {
            return "", 0, fmt.Errorf("Error sending to the rtpproxy on " + self.userv.address.String() + ": " + err.Error())
        }
        retries++
        if s != nil {
            s.Close()
        }
        s, err = net.Dial(self.userv.address.Network(), self.userv.address.String())
        if err != nil {
            time.Sleep(100 * time.Millisecond)
            continue
        }
        _, err = s.Write([]byte(command))
        if err != nil {
            time.Sleep(100 * time.Millisecond)
            continue
        }
        n, err = s.Read(buf)
        if err != nil {
            time.Sleep(100 * time.Millisecond)
            continue
        }
        rval = strings.TrimSpace(string(buf[:n]))
        break
    }
    if s != nil {
        s.Close()
    }
    rtpc_delay, _ := stime.OffsetFromNow()
    return rval, rtpc_delay, nil
}

func (self *_RTPPLWorker) run() {
    for {
        req := <-self.userv.wi
        if req == nil {
            // Shutdown request, relay it further
            self.userv.wi <- nil
            break
        }
        //command, result_callback, callback_parameters = wi
        data, rtpc_delay, err := self.send_raw(req.command, nil)
        if err != nil {
            self.userv.owner.logger.Debug("Error communicating the rtpproxy: " + err.Error())
            data, rtpc_delay = "", -1
        }
        if len(data) == 0 {
            rtpc_delay = -1
        }
        if req.result_callback != nil {
            go req.result_callback(data)
        }
        if rtpc_delay != -1 {
            go self.userv.register_delay(rtpc_delay)
        }
    }
    self.shutdown_chan <- 1
}

type Rtp_proxy_client_stream struct {
    owner       *Rtp_proxy_client_base
    address     net.Addr
    nworkers    int
    workers     []*_RTPPLWorker
    delay_flt   sippy_math.RecFilter
    _is_local    bool
    wi          chan *rtpp_req_stream
}

func NewRtp_proxy_client_stream(owner *Rtp_proxy_client_base, global_config sippy_conf.Config, address net.Addr, opts *Rtp_proxy_opts) (rtp_proxy_transport, error) {
    var err error
    if address == nil {
        address, err = net.ResolveUnixAddr("unix", "/var/run/rtpproxy.sock")
        if err != nil {
            return nil, err
        }
    }
    nworkers := 4
    if opts != nil && opts.Nworkers != nil {
        nworkers = *opts.Nworkers
    }
    self := &Rtp_proxy_client_stream{
        owner       : owner,
        address     : address,
        nworkers    : nworkers,
        workers     : make([]*_RTPPLWorker, nworkers),
        delay_flt   : sippy_math.NewRecFilter(0.95, 0.25),
        wi          : make(chan *rtpp_req_stream, 1000),
    }
    if strings.HasPrefix(address.Network(), "unix") {
        self._is_local = true
    } else {
        self._is_local = false
    }
    //self.wi_available = Condition()
    //self.wi = []
    for i := 0; i < self.nworkers; i++ {
        self.workers[i] = NewRTPPLWorker(self)
    }
    return self, nil
}

func (self *Rtp_proxy_client_stream) is_local() bool {
    return self._is_local
}

func (self *Rtp_proxy_client_stream) send_command(command string, result_callback func(string)) {
    if command[len(command)-1] != '\n' {
        command += "\n"
    }
    self.wi <- &rtpp_req_stream{command, result_callback }
}
/*
    def reconnect(self, address, bind_address = nil):
        self.shutdown()
        self.address = address
        self.workers = []
        for i in range(0, self.nworkers):
            self.workers.append(_RTPPLWorker(self))
        self.delay_flt = recfilter(0.95, 0.25)
*/
func (self *Rtp_proxy_client_stream) shutdown() {
    self.wi <- nil
    for _, rworker := range self.workers {
        <-rworker.shutdown_chan
    }
    self.workers = nil
}

func (self *Rtp_proxy_client_stream) register_delay(rtpc_delay time.Duration) {
    self.delay_flt.Apply(rtpc_delay.Seconds())
}
/*
    def get_rtpc_delay(self):
        return self.delay_flt.lastval

if __name__ == "__main__":
    from twisted.internet import reactor
    def display(*args):
        print args
        reactor.crash()
    r = Rtp_proxy_client_stream({"_sip_address":"1.2.3.4"})
    r.send_command("VF 123456", display, "abcd")
    reactor.run(installSignalHandlers = 1)
    r.shutdown()
*/

func _net_errno(err error) syscall.Errno {
    oe, ok := err.(*net.OpError)
    if ! ok { return 0 }
    errno, ok := oe.Err.(*os.SyscallError)
    if ! ok { return 0 }
    ret, ok := errno.Err.(syscall.Errno)
    if ! ok { return 0 }
    return ret
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package wrangler

import (
	"fmt"
	"time"

	"google.golang.org/protobuf/proto"

	"context"

	"vitess.io/vitess/go/vt/key"
	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/topotools"

	querypb "vitess.io/vitess/go/vt/proto/query"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

// Tablet related methods for wrangler

// InitTablet creates or updates a tablet. If no parent is specified
// in the tablet, and the tablet has a replica type, we will find the
// appropriate parent. If createShardAndKeyspace is true and the
// parent keyspace or shard don't exist, they will be created.  If
// allowUpdate is true, and a tablet with the same ID exists, just update it.
// If a tablet is created as master, and there is already a different
// master in the shard, allowMasterOverride must be set.
func (wr *Wrangler) InitTablet(ctx context.Context, tablet *topodatapb.Tablet, allowMasterOverride, createShardAndKeyspace, allowUpdate bool) error {
	shard, kr, err := topo.ValidateShardName(tablet.Shard)
	if err != nil {
		return err
	}
	tablet.Shard = shard
	tablet.KeyRange = kr

	// get the shard, possibly creating it
	var si *topo.ShardInfo

	if createShardAndKeyspace {
		// create the parent keyspace and shard if needed
		si, err = wr.ts.GetOrCreateShard(ctx, tablet.Keyspace, tablet.Shard)
	} else {
		si, err = wr.ts.GetShard(ctx, tablet.Keyspace, tablet.Shard)
		if topo.IsErrType(err, topo.NoNode) {
			return fmt.Errorf("missing parent shard, use -parent option to create it, or CreateKeyspace / CreateShard")
		}
	}

	// get the shard, checks a couple things
	if err != nil {
		return fmt.Errorf("cannot get (or create) shard %v/%v: %v", tablet.Keyspace, tablet.Shard, err)
	}
	if !key.KeyRangeEqual(si.KeyRange, tablet.KeyRange) {
		return fmt.Errorf("shard %v/%v has a different KeyRange: %v != %v", tablet.Keyspace, tablet.Shard, si.KeyRange, tablet.KeyRange)
	}
	if tablet.Type == topodatapb.TabletType_MASTER && si.HasMaster() && !topoproto.TabletAliasEqual(si.MasterAlias, tablet.Alias) && !allowMasterOverride {
		return fmt.Errorf("creating this tablet would override old master %v in shard %v/%v, use allow_master_override flag", topoproto.TabletAliasString(si.MasterAlias), tablet.Keyspace, tablet.Shard)
	}

	if tablet.Type == topodatapb.TabletType_MASTER {
		// we update master_term_start_time even if the master hasn't changed
		// because that means a new master term with the same master
		tablet.MasterTermStartTime = logutil.TimeToProto(time.Now())
	}

	err = wr.ts.CreateTablet(ctx, tablet)
	if topo.IsErrType(err, topo.NodeExists) && allowUpdate {
		// Try to update then
		oldTablet, err := wr.ts.GetTablet(ctx, tablet.Alias)
		if err != nil {
			return fmt.Errorf("failed reading existing tablet %v: %v", topoproto.TabletAliasString(tablet.Alias), err)
		}

		// Check we have the same keyspace / shard, and if not,
		// require the allowDifferentShard flag.
		if oldTablet.Keyspace != tablet.Keyspace || oldTablet.Shard != tablet.Shard {
			return fmt.Errorf("old tablet has shard %v/%v. Cannot override with shard %v/%v. Delete and re-add tablet if you want to change the tablet's keyspace/shard", oldTablet.Keyspace, oldTablet.Shard, tablet.Keyspace, tablet.Shard)
		}
		oldTablet.Tablet = proto.Clone(tablet).(*topodatapb.Tablet)
		if err := wr.ts.UpdateTablet(ctx, oldTablet); err != nil {
			return fmt.Errorf("failed updating tablet %v: %v", topoproto.TabletAliasString(tablet.Alias), err)
		}
		return nil
	}
	return err
}

// DeleteTablet removes a tablet from a shard.
// - if allowMaster is set, we can Delete a master tablet (and clear
// its record from the Shard record if it was the master).
func (wr *Wrangler) DeleteTablet(ctx context.Context, tabletAlias *topodatapb.TabletAlias, allowMaster bool) (err error) {
	// load the tablet, see if we'll need to rebuild
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	wasMaster, err := wr.isMasterTablet(ctx, ti)
	if err != nil {
		return err
	}

	if wasMaster && !allowMaster {
		return fmt.Errorf("cannot delete tablet %v as it is a master, use allow_master flag", topoproto.TabletAliasString(tabletAlias))
	}

	// update the Shard object if the master was scrapped.
	// we do this before calling DeleteTablet so that the operation can be retried in case of failure.
	if wasMaster {
		// We lock the shard to not conflict with reparent operations.
		ctx, unlock, lockErr := wr.ts.LockShard(ctx, ti.Keyspace, ti.Shard, fmt.Sprintf("DeleteTablet(%v)", topoproto.TabletAliasString(tabletAlias)))
		if lockErr != nil {
			return lockErr
		}
		defer unlock(&err)

		// update the shard record's master
		if _, err := wr.ts.UpdateShardFields(ctx, ti.Keyspace, ti.Shard, func(si *topo.ShardInfo) error {
			if !topoproto.TabletAliasEqual(si.MasterAlias, tabletAlias) {
				wr.Logger().Warningf("Deleting master %v from shard %v/%v but master in Shard object was %v", topoproto.TabletAliasString(tabletAlias), ti.Keyspace, ti.Shard, topoproto.TabletAliasString(si.MasterAlias))
				return topo.NewError(topo.NoUpdateNeeded, si.Keyspace()+"/"+si.ShardName())
			}
			si.MasterAlias = nil
			si.SetMasterTermStartTime(time.Now())
			return nil
		}); err != nil {
			return err
		}
	}

	// remove the record and its replication graph entry
	if err := topotools.DeleteTablet(ctx, wr.ts, ti.Tablet); err != nil {
		return err
	}

	return
}

// ChangeTabletType changes the type of tablet and recomputes all
// necessary derived paths in the serving graph, if necessary.
//
// Note we don't update the master record in the Shard here, as we
// can't ChangeType from and out of master anyway.
func (wr *Wrangler) ChangeTabletType(ctx context.Context, tabletAlias *topodatapb.TabletAlias, tabletType topodatapb.TabletType) error {
	// Load tablet to find endpoint, and keyspace and shard assignment.
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	if !topo.IsTrivialTypeChange(ti.Type, tabletType) {
		return fmt.Errorf("tablet %v type change %v -> %v is not an allowed transition for ChangeTabletType", tabletAlias, ti.Type, tabletType)
	}

	// and ask the tablet to make the change
	return wr.tmc.ChangeType(ctx, ti.Tablet, tabletType)
}

// RefreshTabletState refreshes tablet state
func (wr *Wrangler) RefreshTabletState(ctx context.Context, tabletAlias *topodatapb.TabletAlias) error {
	// Load tablet to find endpoint, and keyspace and shard assignment.
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	// and ask the tablet to refresh itself
	return wr.tmc.RefreshState(ctx, ti.Tablet)
}

// ExecuteFetchAsApp executes a query remotely using the App pool
func (wr *Wrangler) ExecuteFetchAsApp(ctx context.Context, tabletAlias *topodatapb.TabletAlias, usePool bool, query string, maxRows int) (*querypb.QueryResult, error) {
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}
	return wr.tmc.ExecuteFetchAsApp(ctx, ti.Tablet, usePool, []byte(query), maxRows)
}

// ExecuteFetchAsDba executes a query remotely using the DBA pool
func (wr *Wrangler) ExecuteFetchAsDba(ctx context.Context, tabletAlias *topodatapb.TabletAlias, query string, maxRows int, disableBinlogs bool, reloadSchema bool) (*querypb.QueryResult, error) {
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}
	return wr.tmc.ExecuteFetchAsDba(ctx, ti.Tablet, false, []byte(query), maxRows, disableBinlogs, reloadSchema)
}

// VReplicationExec executes a query remotely using the DBA pool
func (wr *Wrangler) VReplicationExec(ctx context.Context, tabletAlias *topodatapb.TabletAlias, query string) (*querypb.QueryResult, error) {
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}
	return wr.tmc.VReplicationExec(ctx, ti.Tablet, query)
}

// GenericVExec executes a query remotely using the DBA pool
func (wr *Wrangler) GenericVExec(ctx context.Context, tabletAlias *topodatapb.TabletAlias, query, workflow, keyspace string) (*querypb.QueryResult, error) {
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}
	return wr.tmc.VExec(ctx, ti.Tablet, query, workflow, keyspace)
}

// isMasterTablet is a shortcut way to determine whether the current tablet
// is a master before we allow its tablet record to be deleted. The canonical
// way to determine the only true master in a shard is to list all the tablets
// and find the one with the highest MasterTermStartTime among the ones that
// claim to be master.
// We err on the side of caution here, i.e. we should never return false for
// a true master tablet, but it is ok to return true for a tablet that isn't
// the true master. This can occur if someone issues a DeleteTablet while
// the system is in transition (a reparenting event is in progress and parts of
// the topo have not yet been updated).
func (wr *Wrangler) isMasterTablet(ctx context.Context, ti *topo.TabletInfo) (bool, error) {
	return topotools.IsPrimaryTablet(ctx, wr.TopoServer(), ti)
}
<commit_msg>revert unnecessary change<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package wrangler

import (
	"fmt"
	"time"

	"google.golang.org/protobuf/proto"

	"context"

	"vitess.io/vitess/go/vt/key"
	"vitess.io/vitess/go/vt/logutil"
	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/topoproto"
	"vitess.io/vitess/go/vt/topotools"

	querypb "vitess.io/vitess/go/vt/proto/query"
	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

// Tablet related methods for wrangler

// InitTablet creates or updates a tablet. If no parent is specified
// in the tablet, and the tablet has a replica type, we will find the
// appropriate parent. If createShardAndKeyspace is true and the
// parent keyspace or shard don't exist, they will be created.  If
// allowUpdate is true, and a tablet with the same ID exists, just update it.
// If a tablet is created as master, and there is already a different
// master in the shard, allowMasterOverride must be set.
func (wr *Wrangler) InitTablet(ctx context.Context, tablet *topodatapb.Tablet, allowMasterOverride, createShardAndKeyspace, allowUpdate bool) error {
	shard, kr, err := topo.ValidateShardName(tablet.Shard)
	if err != nil {
		return err
	}
	tablet.Shard = shard
	tablet.KeyRange = kr

	// get the shard, possibly creating it
	var si *topo.ShardInfo

	if createShardAndKeyspace {
		// create the parent keyspace and shard if needed
		si, err = wr.ts.GetOrCreateShard(ctx, tablet.Keyspace, tablet.Shard)
	} else {
		si, err = wr.ts.GetShard(ctx, tablet.Keyspace, tablet.Shard)
		if topo.IsErrType(err, topo.NoNode) {
			return fmt.Errorf("missing parent shard, use -parent option to create it, or CreateKeyspace / CreateShard")
		}
	}

	// get the shard, checks a couple things
	if err != nil {
		return fmt.Errorf("cannot get (or create) shard %v/%v: %v", tablet.Keyspace, tablet.Shard, err)
	}
	if !key.KeyRangeEqual(si.KeyRange, tablet.KeyRange) {
		return fmt.Errorf("shard %v/%v has a different KeyRange: %v != %v", tablet.Keyspace, tablet.Shard, si.KeyRange, tablet.KeyRange)
	}
	if tablet.Type == topodatapb.TabletType_MASTER && si.HasMaster() && !topoproto.TabletAliasEqual(si.MasterAlias, tablet.Alias) && !allowMasterOverride {
		return fmt.Errorf("creating this tablet would override old master %v in shard %v/%v, use allow_master_override flag", topoproto.TabletAliasString(si.MasterAlias), tablet.Keyspace, tablet.Shard)
	}

	if tablet.Type == topodatapb.TabletType_MASTER {
		// we update master_term_start_time even if the master hasn't changed
		// because that means a new master term with the same master
		tablet.MasterTermStartTime = logutil.TimeToProto(time.Now())
	}

	err = wr.ts.CreateTablet(ctx, tablet)
	if topo.IsErrType(err, topo.NodeExists) && allowUpdate {
		// Try to update then
		oldTablet, err := wr.ts.GetTablet(ctx, tablet.Alias)
		if err != nil {
			return fmt.Errorf("failed reading existing tablet %v: %v", topoproto.TabletAliasString(tablet.Alias), err)
		}

		// Check we have the same keyspace / shard, and if not,
		// require the allowDifferentShard flag.
		if oldTablet.Keyspace != tablet.Keyspace || oldTablet.Shard != tablet.Shard {
			return fmt.Errorf("old tablet has shard %v/%v. Cannot override with shard %v/%v. Delete and re-add tablet if you want to change the tablet's keyspace/shard", oldTablet.Keyspace, oldTablet.Shard, tablet.Keyspace, tablet.Shard)
		}
		oldTablet.Tablet = proto.Clone(tablet).(*topodatapb.Tablet)
		if err := wr.ts.UpdateTablet(ctx, oldTablet); err != nil {
			return fmt.Errorf("failed updating tablet %v: %v", topoproto.TabletAliasString(tablet.Alias), err)
		}
		return nil
	}
	return err
}

// DeleteTablet removes a tablet from a shard.
// - if allowMaster is set, we can Delete a master tablet (and clear
// its record from the Shard record if it was the master).
func (wr *Wrangler) DeleteTablet(ctx context.Context, tabletAlias *topodatapb.TabletAlias, allowMaster bool) (err error) {
	// load the tablet, see if we'll need to rebuild
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	wasMaster, err := wr.isMasterTablet(ctx, ti)
	if err != nil {
		return err
	}

	if wasMaster && !allowMaster {
		return fmt.Errorf("cannot delete tablet %v as it is a master, use allow_master flag", topoproto.TabletAliasString(tabletAlias))
	}

	// update the Shard object if the master was scrapped.
	// we do this before calling DeleteTablet so that the operation can be retried in case of failure.
	if wasMaster {
		// We lock the shard to not conflict with reparent operations.
		ctx, unlock, lockErr := wr.ts.LockShard(ctx, ti.Keyspace, ti.Shard, fmt.Sprintf("DeleteTablet(%v)", topoproto.TabletAliasString(tabletAlias)))
		if lockErr != nil {
			return lockErr
		}
		defer unlock(&err)

		// update the shard record's master
		if _, err := wr.ts.UpdateShardFields(ctx, ti.Keyspace, ti.Shard, func(si *topo.ShardInfo) error {
			if !topoproto.TabletAliasEqual(si.MasterAlias, tabletAlias) {
				wr.Logger().Warningf("Deleting master %v from shard %v/%v but master in Shard object was %v", topoproto.TabletAliasString(tabletAlias), ti.Keyspace, ti.Shard, topoproto.TabletAliasString(si.MasterAlias))
				return topo.NewError(topo.NoUpdateNeeded, si.Keyspace()+"/"+si.ShardName())
			}
			si.MasterAlias = nil
			si.SetMasterTermStartTime(time.Now())
			return nil
		}); err != nil {
			return err
		}
	}

	// remove the record and its replication graph entry
	if err := topotools.DeleteTablet(ctx, wr.ts, ti.Tablet); err != nil {
		return err
	}

	return nil
}

// ChangeTabletType changes the type of tablet and recomputes all
// necessary derived paths in the serving graph, if necessary.
//
// Note we don't update the master record in the Shard here, as we
// can't ChangeType from and out of master anyway.
func (wr *Wrangler) ChangeTabletType(ctx context.Context, tabletAlias *topodatapb.TabletAlias, tabletType topodatapb.TabletType) error {
	// Load tablet to find endpoint, and keyspace and shard assignment.
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	if !topo.IsTrivialTypeChange(ti.Type, tabletType) {
		return fmt.Errorf("tablet %v type change %v -> %v is not an allowed transition for ChangeTabletType", tabletAlias, ti.Type, tabletType)
	}

	// and ask the tablet to make the change
	return wr.tmc.ChangeType(ctx, ti.Tablet, tabletType)
}

// RefreshTabletState refreshes tablet state
func (wr *Wrangler) RefreshTabletState(ctx context.Context, tabletAlias *topodatapb.TabletAlias) error {
	// Load tablet to find endpoint, and keyspace and shard assignment.
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return err
	}

	// and ask the tablet to refresh itself
	return wr.tmc.RefreshState(ctx, ti.Tablet)
}

// ExecuteFetchAsApp executes a query remotely using the App pool
func (wr *Wrangler) ExecuteFetchAsApp(ctx context.Context, tabletAlias *topodatapb.TabletAlias, usePool bool, query string, maxRows int) (*querypb.QueryResult, error) {
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}
	return wr.tmc.ExecuteFetchAsApp(ctx, ti.Tablet, usePool, []byte(query), maxRows)
}

// ExecuteFetchAsDba executes a query remotely using the DBA pool
func (wr *Wrangler) ExecuteFetchAsDba(ctx context.Context, tabletAlias *topodatapb.TabletAlias, query string, maxRows int, disableBinlogs bool, reloadSchema bool) (*querypb.QueryResult, error) {
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}
	return wr.tmc.ExecuteFetchAsDba(ctx, ti.Tablet, false, []byte(query), maxRows, disableBinlogs, reloadSchema)
}

// VReplicationExec executes a query remotely using the DBA pool
func (wr *Wrangler) VReplicationExec(ctx context.Context, tabletAlias *topodatapb.TabletAlias, query string) (*querypb.QueryResult, error) {
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}
	return wr.tmc.VReplicationExec(ctx, ti.Tablet, query)
}

// GenericVExec executes a query remotely using the DBA pool
func (wr *Wrangler) GenericVExec(ctx context.Context, tabletAlias *topodatapb.TabletAlias, query, workflow, keyspace string) (*querypb.QueryResult, error) {
	ti, err := wr.ts.GetTablet(ctx, tabletAlias)
	if err != nil {
		return nil, err
	}
	return wr.tmc.VExec(ctx, ti.Tablet, query, workflow, keyspace)
}

// isMasterTablet is a shortcut way to determine whether the current tablet
// is a master before we allow its tablet record to be deleted. The canonical
// way to determine the only true master in a shard is to list all the tablets
// and find the one with the highest MasterTermStartTime among the ones that
// claim to be master.
// We err on the side of caution here, i.e. we should never return false for
// a true master tablet, but it is ok to return true for a tablet that isn't
// the true master. This can occur if someone issues a DeleteTablet while
// the system is in transition (a reparenting event is in progress and parts of
// the topo have not yet been updated).
func (wr *Wrangler) isMasterTablet(ctx context.Context, ti *topo.TabletInfo) (bool, error) {
	return topotools.IsPrimaryTablet(ctx, wr.TopoServer(), ti)
}
<|endoftext|>


<commit_before>package gocbcore

import (
	"encoding/binary"
)

// Represents the state of a particular cluster snapshot.
type SnapshotState uint32

// Returns whether this snapshot is available in memory.
func (s SnapshotState) HasInMemory() bool {
	return uint32(s)&1 != 0
}

// Returns whether this snapshot is available on disk.
func (s SnapshotState) HasOnDisk() bool {
	return uint32(s)&2 != 0
}

// Represents a single entry in the server failover log.
type FailoverEntry struct {
	VbUuid VbUuid
	SeqNo  SeqNo
}

type StreamObserver interface {
	SnapshotMarker(startSeqNo, endSeqNo uint64, vbId uint16, snapshotType SnapshotState)
	Mutation(seqNo, revNo uint64, flags, expiry, lockTime uint32, cas uint64, datatype uint8, vbId uint16, key, value []byte)
	Deletion(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	Expiration(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	End(vbId uint16, err error)
}

type OpenStreamCallback func([]FailoverEntry, error)
type CloseStreamCallback func(error)
type GetFailoverLogCallback func([]FailoverEntry, error)
type GetVBucketSeqnosCallback func(uint16, SeqNo, error)

// **INTERNAL**
// Opens a DCP stream for a particular VBucket.
func (c *Agent) OpenStream(vbId uint16, vbUuid VbUuid, startSeqNo, endSeqNo, snapStartSeqNo, snapEndSeqNo SeqNo, evtHandler StreamObserver, cb OpenStreamCallback) (PendingOp, error) {
	var req *memdQRequest
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			// All client errors are handled by the StreamObserver
			cb(nil, err)
			return
		}

		if resp.Magic == ResMagic {
			// This is the response to the open stream request.

			numEntries := len(resp.Value) / 16
			entries := make([]FailoverEntry, numEntries)
			for i := 0; i < numEntries; i++ {
				entries[i] = FailoverEntry{
					VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
					SeqNo:  SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
				}
			}

			cb(entries, nil)
			return
		}

		// This is one of the stream events
		switch resp.Opcode {
		case CmdDcpSnapshotMarker:
			vbId := uint16(resp.Status)
			newStartSeqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			newEndSeqNo := binary.BigEndian.Uint64(resp.Extras[8:])
			snapshotType := binary.BigEndian.Uint32(resp.Extras[16:])
			evtHandler.SnapshotMarker(newStartSeqNo, newEndSeqNo, vbId, SnapshotState(snapshotType))
		case CmdDcpMutation:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			flags := binary.BigEndian.Uint32(resp.Extras[16:])
			expiry := binary.BigEndian.Uint32(resp.Extras[20:])
			lockTime := binary.BigEndian.Uint32(resp.Extras[24:])
			evtHandler.Mutation(seqNo, revNo, flags, expiry, lockTime, resp.Cas, resp.Datatype, vbId, resp.Key, resp.Value)
		case CmdDcpDeletion:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Deletion(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpExpiration:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Expiration(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpStreamEnd:
			vbId := uint16(resp.Status)
			code := StreamEndStatus(binary.BigEndian.Uint32(resp.Extras[0:]))
			evtHandler.End(vbId, getStreamEndError(code))
			req.Cancel()
		}
	}

	extraBuf := make([]byte, 48)
	binary.BigEndian.PutUint32(extraBuf[0:], 0)
	binary.BigEndian.PutUint32(extraBuf[4:], 0)
	binary.BigEndian.PutUint64(extraBuf[8:], uint64(startSeqNo))
	binary.BigEndian.PutUint64(extraBuf[16:], uint64(endSeqNo))
	binary.BigEndian.PutUint64(extraBuf[24:], uint64(vbUuid))
	binary.BigEndian.PutUint64(extraBuf[32:], uint64(snapStartSeqNo))
	binary.BigEndian.PutUint64(extraBuf[40:], uint64(snapEndSeqNo))

	req = &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpStreamReq,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

// **INTERNAL**
// Shuts down an open stream for the specified VBucket.
func (c *Agent) CloseStream(vbId uint16, cb CloseStreamCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		cb(err)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpCloseStream,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: false,
	}
	return c.dispatchOp(req)
}

// **INTERNAL**
// Retrieves the failover log for a particular VBucket.  This is used
// to resume an interrupted stream after a node failover has occured.
func (c *Agent) GetFailoverLog(vbId uint16, cb GetFailoverLogCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			cb(nil, err)
			return
		}

		numEntries := len(resp.Value) / 16
		entries := make([]FailoverEntry, numEntries)
		for i := 0; i < numEntries; i++ {
			entries[i] = FailoverEntry{
				VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
				SeqNo:  SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
			}
		}
		cb(entries, nil)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpGetFailoverLog,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: false,
	}
	return c.dispatchOp(req)
}

// **INTERNAL**
// Returns the last checkpoint for a particular VBucket.  This is useful
// for starting a DCP stream from wherever the server currently is.
func (c *Agent) GetVbucketSeqnos(serverIdx int, cb GetVBucketSeqnosCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			cb(0, 0, err)
			return
		}

		vbs := len(resp.Value) / 10
		for i := 0; i < vbs; i++ {
			vbid := binary.BigEndian.Uint16(resp.Value[i*10:])
			seqNo := SeqNo(binary.BigEndian.Uint64(resp.Value[i*10+2:]))
			cb(vbid, seqNo, nil)
		}
	}

	extraBuf := make([]byte, 4)
	binary.BigEndian.PutUint32(extraBuf[0:], uint32(VBucketStateActive))

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdGetAllVBSeqnos,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  0,
		},
		Callback:   handler,
		ReplicaIdx: serverIdx,
		Persistent: false,
	}

	return c.dispatchOp(req)
}
<commit_msg>Server indexes should be negative for a broadcast operation<commit_after>package gocbcore

import (
	"encoding/binary"
)

// Represents the state of a particular cluster snapshot.
type SnapshotState uint32

// Returns whether this snapshot is available in memory.
func (s SnapshotState) HasInMemory() bool {
	return uint32(s)&1 != 0
}

// Returns whether this snapshot is available on disk.
func (s SnapshotState) HasOnDisk() bool {
	return uint32(s)&2 != 0
}

// Represents a single entry in the server failover log.
type FailoverEntry struct {
	VbUuid VbUuid
	SeqNo  SeqNo
}

type StreamObserver interface {
	SnapshotMarker(startSeqNo, endSeqNo uint64, vbId uint16, snapshotType SnapshotState)
	Mutation(seqNo, revNo uint64, flags, expiry, lockTime uint32, cas uint64, datatype uint8, vbId uint16, key, value []byte)
	Deletion(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	Expiration(seqNo, revNo, cas uint64, vbId uint16, key []byte)
	End(vbId uint16, err error)
}

type OpenStreamCallback func([]FailoverEntry, error)
type CloseStreamCallback func(error)
type GetFailoverLogCallback func([]FailoverEntry, error)
type GetVBucketSeqnosCallback func(uint16, SeqNo, error)

// **INTERNAL**
// Opens a DCP stream for a particular VBucket.
func (c *Agent) OpenStream(vbId uint16, vbUuid VbUuid, startSeqNo, endSeqNo, snapStartSeqNo, snapEndSeqNo SeqNo, evtHandler StreamObserver, cb OpenStreamCallback) (PendingOp, error) {
	var req *memdQRequest
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			// All client errors are handled by the StreamObserver
			cb(nil, err)
			return
		}

		if resp.Magic == ResMagic {
			// This is the response to the open stream request.

			numEntries := len(resp.Value) / 16
			entries := make([]FailoverEntry, numEntries)
			for i := 0; i < numEntries; i++ {
				entries[i] = FailoverEntry{
					VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
					SeqNo:  SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
				}
			}

			cb(entries, nil)
			return
		}

		// This is one of the stream events
		switch resp.Opcode {
		case CmdDcpSnapshotMarker:
			vbId := uint16(resp.Status)
			newStartSeqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			newEndSeqNo := binary.BigEndian.Uint64(resp.Extras[8:])
			snapshotType := binary.BigEndian.Uint32(resp.Extras[16:])
			evtHandler.SnapshotMarker(newStartSeqNo, newEndSeqNo, vbId, SnapshotState(snapshotType))
		case CmdDcpMutation:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			flags := binary.BigEndian.Uint32(resp.Extras[16:])
			expiry := binary.BigEndian.Uint32(resp.Extras[20:])
			lockTime := binary.BigEndian.Uint32(resp.Extras[24:])
			evtHandler.Mutation(seqNo, revNo, flags, expiry, lockTime, resp.Cas, resp.Datatype, vbId, resp.Key, resp.Value)
		case CmdDcpDeletion:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Deletion(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpExpiration:
			vbId := uint16(resp.Status)
			seqNo := binary.BigEndian.Uint64(resp.Extras[0:])
			revNo := binary.BigEndian.Uint64(resp.Extras[8:])
			evtHandler.Expiration(seqNo, revNo, resp.Cas, vbId, resp.Key)
		case CmdDcpStreamEnd:
			vbId := uint16(resp.Status)
			code := StreamEndStatus(binary.BigEndian.Uint32(resp.Extras[0:]))
			evtHandler.End(vbId, getStreamEndError(code))
			req.Cancel()
		}
	}

	extraBuf := make([]byte, 48)
	binary.BigEndian.PutUint32(extraBuf[0:], 0)
	binary.BigEndian.PutUint32(extraBuf[4:], 0)
	binary.BigEndian.PutUint64(extraBuf[8:], uint64(startSeqNo))
	binary.BigEndian.PutUint64(extraBuf[16:], uint64(endSeqNo))
	binary.BigEndian.PutUint64(extraBuf[24:], uint64(vbUuid))
	binary.BigEndian.PutUint64(extraBuf[32:], uint64(snapStartSeqNo))
	binary.BigEndian.PutUint64(extraBuf[40:], uint64(snapEndSeqNo))

	req = &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpStreamReq,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: true,
	}
	return c.dispatchOp(req)
}

// **INTERNAL**
// Shuts down an open stream for the specified VBucket.
func (c *Agent) CloseStream(vbId uint16, cb CloseStreamCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		cb(err)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpCloseStream,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: false,
	}
	return c.dispatchOp(req)
}

// **INTERNAL**
// Retrieves the failover log for a particular VBucket.  This is used
// to resume an interrupted stream after a node failover has occured.
func (c *Agent) GetFailoverLog(vbId uint16, cb GetFailoverLogCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			cb(nil, err)
			return
		}

		numEntries := len(resp.Value) / 16
		entries := make([]FailoverEntry, numEntries)
		for i := 0; i < numEntries; i++ {
			entries[i] = FailoverEntry{
				VbUuid: VbUuid(binary.BigEndian.Uint64(resp.Value[i*16+0:])),
				SeqNo:  SeqNo(binary.BigEndian.Uint64(resp.Value[i*16+8:])),
			}
		}
		cb(entries, nil)
	}

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdDcpGetFailoverLog,
			Datatype: 0,
			Cas:      0,
			Extras:   nil,
			Key:      nil,
			Value:    nil,
			Vbucket:  vbId,
		},
		Callback:   handler,
		ReplicaIdx: 0,
		Persistent: false,
	}
	return c.dispatchOp(req)
}

// **INTERNAL**
// Returns the last checkpoint for a particular VBucket.  This is useful
// for starting a DCP stream from wherever the server currently is.
func (c *Agent) GetVbucketSeqnos(serverIdx int, cb GetVBucketSeqnosCallback) (PendingOp, error) {
	handler := func(resp *memdResponse, err error) {
		if err != nil {
			cb(0, 0, err)
			return
		}

		vbs := len(resp.Value) / 10
		for i := 0; i < vbs; i++ {
			vbid := binary.BigEndian.Uint16(resp.Value[i*10:])
			seqNo := SeqNo(binary.BigEndian.Uint64(resp.Value[i*10+2:]))
			cb(vbid, seqNo, nil)
		}
	}

	extraBuf := make([]byte, 4)
	binary.BigEndian.PutUint32(extraBuf[0:], uint32(VBucketStateActive))

	req := &memdQRequest{
		memdRequest: memdRequest{
			Magic:    ReqMagic,
			Opcode:   CmdGetAllVBSeqnos,
			Datatype: 0,
			Cas:      0,
			Extras:   extraBuf,
			Key:      nil,
			Value:    nil,
			Vbucket:  0,
		},
		Callback:   handler,
		ReplicaIdx: -serverIdx,
		Persistent: false,
	}

	return c.dispatchOp(req)
}
<|endoftext|>


<commit_before>// Package bootstrap implements the bootstrapping logic: generation of a .go file to
// launch the actual generator and launching the generator itself.
//
// The package may be preferred to a command-line utility if generating the serializers
// from golang code is required.
package bootstrap

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
)

const genPackage = "github.com/mailru/easyjson/gen"
const pkgWriter = "github.com/mailru/easyjson/jwriter"
const pkgLexer = "github.com/mailru/easyjson/jlexer"

type Generator struct {
	PkgPath, PkgName string
	Types            []string

	NoStdMarshalers       bool
	SnakeCase             bool
	LowerCamelCase        bool
	OmitEmpty             bool
	DisallowUnknownFields bool

	OutName   string
	BuildTags string

	StubsOnly  bool
	LeaveTemps bool
	NoFormat   bool
}

// writeStub outputs an initial stubs for marshalers/unmarshalers so that the package
// using marshalers/unmarshales compiles correctly for boostrapping code.
func (g *Generator) writeStub() error {
	f, err := os.Create(g.OutName)
	if err != nil {
		return err
	}
	defer f.Close()

	if g.BuildTags != "" {
		fmt.Fprintln(f, "// +build ", g.BuildTags)
		fmt.Fprintln(f)
	}
	fmt.Fprintln(f, "// TEMPORARY AUTOGENERATED FILE: easyjson stub code to make the package")
	fmt.Fprintln(f, "// compilable during generation.")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "package ", g.PkgName)

	if len(g.Types) > 0 {
		fmt.Fprintln(f)
		fmt.Fprintln(f, "import (")
		fmt.Fprintln(f, `  "`+pkgWriter+`"`)
		fmt.Fprintln(f, `  "`+pkgLexer+`"`)
		fmt.Fprintln(f, ")")
	}

	sort.Strings(g.Types)
	for _, t := range g.Types {
		fmt.Fprintln(f)
		if !g.NoStdMarshalers {
			fmt.Fprintln(f, "func (", t, ") MarshalJSON() ([]byte, error) { return nil, nil }")
			fmt.Fprintln(f, "func (*", t, ") UnmarshalJSON([]byte) error { return nil }")
		}

		fmt.Fprintln(f, "func (", t, ") MarshalEasyJSON(w *jwriter.Writer) {}")
		fmt.Fprintln(f, "func (*", t, ") UnmarshalEasyJSON(l *jlexer.Lexer) {}")
		fmt.Fprintln(f)
		fmt.Fprintln(f, "type EasyJSON_exporter_"+t+" *"+t)
	}
	return nil
}

// writeMain creates a .go file that launches the generator if 'go run'.
func (g *Generator) writeMain() (path string, err error) {
	f, err := ioutil.TempFile(filepath.Dir(g.OutName), "easyjson-bootstrap")
	if err != nil {
		return "", err
	}

	fmt.Fprintln(f, "// +build ignore")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "// TEMPORARY AUTOGENERATED FILE: easyjson bootstapping code to launch")
	fmt.Fprintln(f, "// the actual generator.")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "package main")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "import (")
	fmt.Fprintln(f, `  "fmt"`)
	fmt.Fprintln(f, `  "os"`)
	fmt.Fprintln(f)
	fmt.Fprintf(f, "  %q\n", genPackage)
	if len(g.Types) > 0 {
		fmt.Fprintln(f)
		fmt.Fprintf(f, "  pkg %q\n", g.PkgPath)
	}
	fmt.Fprintln(f, ")")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "func main() {")
	fmt.Fprintf(f, "  g := gen.NewGenerator(%q)\n", filepath.Base(g.OutName))
	fmt.Fprintf(f, "  g.SetPkg(%q, %q)\n", g.PkgName, g.PkgPath)
	if g.BuildTags != "" {
		fmt.Fprintf(f, "  g.SetBuildTags(%q)\n", g.BuildTags)
	}
	if g.SnakeCase {
		fmt.Fprintln(f, "  g.UseSnakeCase()")
	}
	if g.LowerCamelCase {
		fmt.Fprintln(f, "  g.UseLowerCamelCase()")
	}
	if g.OmitEmpty {
		fmt.Fprintln(f, "  g.OmitEmpty()")
	}
	if g.NoStdMarshalers {
		fmt.Fprintln(f, "  g.NoStdMarshalers()")
	}
	if g.DisallowUnknownFields {
		fmt.Fprintln(f, "  g.DisallowUnknownFields()")
	}

	sort.Strings(g.Types)
	for _, v := range g.Types {
		fmt.Fprintln(f, "  g.Add(pkg.EasyJSON_exporter_"+v+"(nil))")
	}

	fmt.Fprintln(f, "  if err := g.Run(os.Stdout); err != nil {")
	fmt.Fprintln(f, "    fmt.Fprintln(os.Stderr, err)")
	fmt.Fprintln(f, "    os.Exit(1)")
	fmt.Fprintln(f, "  }")
	fmt.Fprintln(f, "}")

	src := f.Name()
	if err := f.Close(); err != nil {
		return src, err
	}

	dest := src + ".go"
	return dest, os.Rename(src, dest)
}

func (g *Generator) Run() error {
	if err := g.writeStub(); err != nil {
		return err
	}
	if g.StubsOnly {
		return nil
	}

	path, err := g.writeMain()
	if err != nil {
		return err
	}
	if !g.LeaveTemps {
		defer os.Remove(path)
	}

	f, err := os.Create(g.OutName + ".tmp")
	if err != nil {
		return err
	}
	if !g.LeaveTemps {
		defer os.Remove(f.Name()) // will not remove after rename
	}

	cmd := exec.Command("go", "run", "-tags", g.BuildTags, path)
	cmd.Stdout = f
	cmd.Stderr = os.Stderr
	cmd.Dir = filepath.Dir(path)
	if err = cmd.Run(); err != nil {
		return err
	}

	f.Close()

	if !g.NoFormat {
		cmd = exec.Command("gofmt", "-w", f.Name())
		cmd.Stderr = os.Stderr
		cmd.Stdout = os.Stdout

		if err = cmd.Run(); err != nil {
			return err
		}
	}

	return os.Rename(f.Name(), g.OutName)
}
<commit_msg>relative paths fix<commit_after>// Package bootstrap implements the bootstrapping logic: generation of a .go file to
// launch the actual generator and launching the generator itself.
//
// The package may be preferred to a command-line utility if generating the serializers
// from golang code is required.
package bootstrap

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
)

const genPackage = "github.com/mailru/easyjson/gen"
const pkgWriter = "github.com/mailru/easyjson/jwriter"
const pkgLexer = "github.com/mailru/easyjson/jlexer"

type Generator struct {
	PkgPath, PkgName string
	Types            []string

	NoStdMarshalers       bool
	SnakeCase             bool
	LowerCamelCase        bool
	OmitEmpty             bool
	DisallowUnknownFields bool

	OutName   string
	BuildTags string

	StubsOnly  bool
	LeaveTemps bool
	NoFormat   bool
}

// writeStub outputs an initial stubs for marshalers/unmarshalers so that the package
// using marshalers/unmarshales compiles correctly for boostrapping code.
func (g *Generator) writeStub() error {
	f, err := os.Create(g.OutName)
	if err != nil {
		return err
	}
	defer f.Close()

	if g.BuildTags != "" {
		fmt.Fprintln(f, "// +build ", g.BuildTags)
		fmt.Fprintln(f)
	}
	fmt.Fprintln(f, "// TEMPORARY AUTOGENERATED FILE: easyjson stub code to make the package")
	fmt.Fprintln(f, "// compilable during generation.")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "package ", g.PkgName)

	if len(g.Types) > 0 {
		fmt.Fprintln(f)
		fmt.Fprintln(f, "import (")
		fmt.Fprintln(f, `  "`+pkgWriter+`"`)
		fmt.Fprintln(f, `  "`+pkgLexer+`"`)
		fmt.Fprintln(f, ")")
	}

	sort.Strings(g.Types)
	for _, t := range g.Types {
		fmt.Fprintln(f)
		if !g.NoStdMarshalers {
			fmt.Fprintln(f, "func (", t, ") MarshalJSON() ([]byte, error) { return nil, nil }")
			fmt.Fprintln(f, "func (*", t, ") UnmarshalJSON([]byte) error { return nil }")
		}

		fmt.Fprintln(f, "func (", t, ") MarshalEasyJSON(w *jwriter.Writer) {}")
		fmt.Fprintln(f, "func (*", t, ") UnmarshalEasyJSON(l *jlexer.Lexer) {}")
		fmt.Fprintln(f)
		fmt.Fprintln(f, "type EasyJSON_exporter_"+t+" *"+t)
	}
	return nil
}

// writeMain creates a .go file that launches the generator if 'go run'.
func (g *Generator) writeMain() (path string, err error) {
	f, err := ioutil.TempFile(filepath.Dir(g.OutName), "easyjson-bootstrap")
	if err != nil {
		return "", err
	}

	fmt.Fprintln(f, "// +build ignore")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "// TEMPORARY AUTOGENERATED FILE: easyjson bootstapping code to launch")
	fmt.Fprintln(f, "// the actual generator.")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "package main")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "import (")
	fmt.Fprintln(f, `  "fmt"`)
	fmt.Fprintln(f, `  "os"`)
	fmt.Fprintln(f)
	fmt.Fprintf(f, "  %q\n", genPackage)
	if len(g.Types) > 0 {
		fmt.Fprintln(f)
		fmt.Fprintf(f, "  pkg %q\n", g.PkgPath)
	}
	fmt.Fprintln(f, ")")
	fmt.Fprintln(f)
	fmt.Fprintln(f, "func main() {")
	fmt.Fprintf(f, "  g := gen.NewGenerator(%q)\n", filepath.Base(g.OutName))
	fmt.Fprintf(f, "  g.SetPkg(%q, %q)\n", g.PkgName, g.PkgPath)
	if g.BuildTags != "" {
		fmt.Fprintf(f, "  g.SetBuildTags(%q)\n", g.BuildTags)
	}
	if g.SnakeCase {
		fmt.Fprintln(f, "  g.UseSnakeCase()")
	}
	if g.LowerCamelCase {
		fmt.Fprintln(f, "  g.UseLowerCamelCase()")
	}
	if g.OmitEmpty {
		fmt.Fprintln(f, "  g.OmitEmpty()")
	}
	if g.NoStdMarshalers {
		fmt.Fprintln(f, "  g.NoStdMarshalers()")
	}
	if g.DisallowUnknownFields {
		fmt.Fprintln(f, "  g.DisallowUnknownFields()")
	}

	sort.Strings(g.Types)
	for _, v := range g.Types {
		fmt.Fprintln(f, "  g.Add(pkg.EasyJSON_exporter_"+v+"(nil))")
	}

	fmt.Fprintln(f, "  if err := g.Run(os.Stdout); err != nil {")
	fmt.Fprintln(f, "    fmt.Fprintln(os.Stderr, err)")
	fmt.Fprintln(f, "    os.Exit(1)")
	fmt.Fprintln(f, "  }")
	fmt.Fprintln(f, "}")

	src := f.Name()
	if err := f.Close(); err != nil {
		return src, err
	}

	dest := src + ".go"
	return dest, os.Rename(src, dest)
}

func (g *Generator) Run() error {
	if err := g.writeStub(); err != nil {
		return err
	}
	if g.StubsOnly {
		return nil
	}

	path, err := g.writeMain()
	if err != nil {
		return err
	}
	if !g.LeaveTemps {
		defer os.Remove(path)
	}

	f, err := os.Create(g.OutName + ".tmp")
	if err != nil {
		return err
	}
	if !g.LeaveTemps {
		defer os.Remove(f.Name()) // will not remove after rename
	}

	cmd := exec.Command("go", "run", "-tags", g.BuildTags, filepath.Base(path))
	cmd.Stdout = f
	cmd.Stderr = os.Stderr
	cmd.Dir = filepath.Dir(path)
	if err = cmd.Run(); err != nil {
		return err
	}

	f.Close()

	if !g.NoFormat {
		cmd = exec.Command("gofmt", "-w", f.Name())
		cmd.Stderr = os.Stderr
		cmd.Stdout = os.Stdout

		if err = cmd.Run(); err != nil {
			return err
		}
	}

	return os.Rename(f.Name(), g.OutName)
}
<|endoftext|>


<commit_before>/*
Package supervisor provides supervisor trees for Go applications.

This package is a clean reimplementation of github.com/thejerf/suture, aiming
to be more Go idiomatic, thus less Erlang-like.

It is built on top of context package, with all of its advantages, namely the
possibility trickle down context-related values and cancellation signals.

TheJerf's blog post about Suture is a very good and helpful read to understand
how this package has been implemented.

http://www.jerf.org/iri/post/2930
*/
package supervisor // import "cirello.io/supervisor"

import (
	"fmt"
	"log"
	"math"
	"sync"
	"time"

	"golang.org/x/net/context"
)

// Service is the public interface expected by a Supervisor.
//
// This will be internally named after the result of fmt.Stringer, if available.
// Otherwise it will going to use an internal representation for the service
// name.
type Service interface {
	// Serve is called by a Supervisor to start the service. It expects the
	// service to honor the passed context and its lifetime. Observe
	// <-ctx.Done() and ctx.Err(). If the service is stopped by anything
	// but the Supervisor, it will get started again. Be careful with shared
	// state among restarts.
	Serve(ctx context.Context)
}

// Supervisor is the basic datastructure responsible for offering a supervisor
// tree. It implements Service, therefore it can be nested if necessary. When
// passing the Supervisor around, remind to do it as reference (&supervisor).
type Supervisor struct {
	// Name for this supervisor tree, used for logging.
	Name string

	// FailureDecay is the timespan on which the current failure count will
	// be halved.
	FailureDecay float64

	// FailureThreshold is the maximum accepted number of failures, after
	// decay adjustment, that shall trigger the back-off wait.
	FailureThreshold float64

	// Backoff is the wait duration when hit threshold.
	Backoff time.Duration

	// Log is a replaceable function used for overall logging
	Log func(interface{})

	ready sync.Once

	addedService    chan struct{}
	startedServices chan struct{}

	runningServices sync.WaitGroup

	servicesMu sync.Mutex
	services   map[string]Service

	cancelationsMu sync.Mutex
	cancelations   map[string]context.CancelFunc

	backoffMu sync.Mutex
	backoff   map[string]*backoff

	runningMu sync.Mutex
}

func (s *Supervisor) String() string {
	return s.Name
}

func (s *Supervisor) prepare() {
	s.ready.Do(func() {
		if s.Name == "" {
			s.Name = "supervisor"
		}
		s.addedService = make(chan struct{}, 1)
		s.backoff = make(map[string]*backoff)
		s.cancelations = make(map[string]context.CancelFunc)
		s.services = make(map[string]Service)
		s.startedServices = make(chan struct{}, 1)

		if s.Log == nil {
			s.Log = func(msg interface{}) {
				log.Printf("%s: %v", s.Name, msg)
			}
		}
		if s.FailureDecay == 0 {
			s.FailureDecay = 30
		}
		if s.FailureThreshold == 0 {
			s.FailureThreshold = 5
		}
		if s.Backoff == 0 {
			s.Backoff = 15 * time.Second
		}
	})
}

// Add inserts into the Supervisor tree a new service. If the Supervisor is
// already started, it will start it automatically.
func (s *Supervisor) Add(service Service) {
	s.prepare()

	name := fmt.Sprintf("%s", service)

	s.servicesMu.Lock()
	s.backoffMu.Lock()
	s.backoff[name] = &backoff{}
	s.services[name] = service
	s.backoffMu.Unlock()
	s.servicesMu.Unlock()

	select {
	case s.addedService <- struct{}{}:
	default:
	}
}

// Remove stops the service in the Supervisor tree and remove from it.
func (s *Supervisor) Remove(name string) {
	s.prepare()

	s.servicesMu.Lock()
	defer s.servicesMu.Unlock()
	if _, ok := s.services[name]; !ok {
		return
	}

	s.cancelationsMu.Lock()
	defer s.cancelationsMu.Unlock()
	if c, ok := s.cancelations[name]; ok {
		delete(s.cancelations, name)
		c()
	}
}

// Serve starts the Supervisor tree. It can be started only once at a time. If
// stopped (canceled), it can be restarted. In case of concurrent calls, it will
// hang until the current call is completed.
func (s *Supervisor) Serve(ctx context.Context) {
	s.prepare()

	s.runningMu.Lock()
	s.serve(ctx)
	s.runningMu.Unlock()
}

// Services return a list of services
func (s *Supervisor) Services() map[string]Service {
	svclist := make(map[string]Service)
	s.servicesMu.Lock()
	for k, v := range s.services {
		svclist[k] = v
	}
	s.servicesMu.Unlock()
	return svclist
}

// Cancelations return a list of services names and their cancellation calls
func (s *Supervisor) Cancelations() map[string]context.CancelFunc {
	svclist := make(map[string]context.CancelFunc)
	s.cancelationsMu.Lock()
	for k, v := range s.cancelations {
		svclist[k] = v
	}
	s.cancelationsMu.Unlock()
	return svclist
}

func (s *Supervisor) startServices(ctx context.Context) {
	s.startAllServices(ctx)
	select {
	case s.startedServices <- struct{}{}:
	default:
	}
}

func (s *Supervisor) serve(ctx context.Context) {
	select {
	case <-s.addedService:
	default:
	}
	s.startServices(ctx)

	go func(ctx context.Context) {
		for {
			select {
			case <-s.addedService:
				s.startServices(ctx)

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

	s.runningServices.Wait()

	s.cancelationsMu.Lock()
	s.cancelations = make(map[string]context.CancelFunc)
	s.cancelationsMu.Unlock()
	return
}

func (s *Supervisor) startAllServices(ctx context.Context) {
	s.servicesMu.Lock()
	defer s.servicesMu.Unlock()

	var wg sync.WaitGroup

	for name, svc := range s.services {
		s.cancelationsMu.Lock()
		if _, ok := s.cancelations[name]; ok {
			s.cancelationsMu.Unlock()
			continue
		}
		s.cancelations[name] = nil
		s.cancelationsMu.Unlock()

		wg.Add(1)
		go func(name string, svc Service) {
			s.runningServices.Add(1)
			wg.Done()
			for {
				retry := func() (retry bool) {
					defer func() {
						if r := recover(); r != nil {
							s.Log(fmt.Sprintf("trapped panic: %v", r))
							retry = true
						}
					}()

					c, cancel := context.WithCancel(ctx)
					s.cancelationsMu.Lock()
					s.cancelations[name] = cancel
					s.cancelationsMu.Unlock()
					svc.Serve(c)

					select {
					case <-ctx.Done():
						return false
					default:
						return true
					}
				}()
				if retry {
					s.Log(fmt.Sprintf("restarting %s", name))
					s.backoffMu.Lock()
					b := s.backoff[name]
					s.backoffMu.Unlock()
					b.wait(s.FailureDecay, s.FailureThreshold, s.Backoff, func(msg interface{}) {
						s.Log(fmt.Sprintf("backoff %s: %v", name, msg))
					})
					continue
				}
				break
			}
			s.runningServices.Done()
		}(name, svc)
	}

	wg.Wait()
}

type backoff struct {
	lastfail time.Time
	failures float64
}

func (b *backoff) wait(failureDecay float64, threshold float64, backoffDur time.Duration, log func(str interface{})) {
	if b.lastfail.IsZero() {
		b.lastfail = time.Now()
		b.failures = 1.0
		return
	}

	b.failures++
	intervals := time.Since(b.lastfail).Seconds() / failureDecay
	b.failures = b.failures*math.Pow(.5, intervals) + 1

	if b.failures > threshold {
		log(backoffDur)
		time.Sleep(backoffDur)
	}
}
<commit_msg>Revert 9f2f9e17b1ddf81de29c196458b9383fb635eafa<commit_after>/*
Package supervisor provides supervisor trees for Go applications.

This package is a clean reimplementation of github.com/thejerf/suture, aiming
to be more Go idiomatic, thus less Erlang-like.

It is built on top of context package, with all of its advantages, namely the
possibility trickle down context-related values and cancellation signals.

TheJerf's blog post about Suture is a very good and helpful read to understand
how this package has been implemented.

http://www.jerf.org/iri/post/2930
*/
package supervisor // import "cirello.io/supervisor"

import (
	"fmt"
	"log"
	"math"
	"sync"
	"time"

	"golang.org/x/net/context"
)

// Service is the public interface expected by a Supervisor.
//
// This will be internally named after the result of fmt.Stringer, if available.
// Otherwise it will going to use an internal representation for the service
// name.
type Service interface {
	// Serve is called by a Supervisor to start the service. It expects the
	// service to honor the passed context and its lifetime. Observe
	// <-ctx.Done() and ctx.Err(). If the service is stopped by anything
	// but the Supervisor, it will get started again. Be careful with shared
	// state among restarts.
	Serve(ctx context.Context)
}

// Supervisor is the basic datastructure responsible for offering a supervisor
// tree. It implements Service, therefore it can be nested if necessary. When
// passing the Supervisor around, remind to do it as reference (&supervisor).
type Supervisor struct {
	// Name for this supervisor tree, used for logging.
	Name string

	// FailureDecay is the timespan on which the current failure count will
	// be halved.
	FailureDecay float64

	// FailureThreshold is the maximum accepted number of failures, after
	// decay adjustment, that shall trigger the back-off wait.
	FailureThreshold float64

	// Backoff is the wait duration when hit threshold.
	Backoff time.Duration

	// Log is a replaceable function used for overall logging
	Log func(interface{})

	ready sync.Once

	addedService    chan struct{}
	startedServices chan struct{}

	runningServices sync.WaitGroup

	servicesMu sync.Mutex
	services   map[string]Service

	cancelationsMu sync.Mutex
	cancelations   map[string]context.CancelFunc

	backoffMu sync.Mutex
	backoff   map[string]*backoff

	runningMu sync.Mutex
}

func (s *Supervisor) String() string {
	return s.Name
}

func (s *Supervisor) prepare() {
	s.ready.Do(func() {
		if s.Name == "" {
			s.Name = "supervisor"
		}
		s.addedService = make(chan struct{}, 1)
		s.backoff = make(map[string]*backoff)
		s.cancelations = make(map[string]context.CancelFunc)
		s.services = make(map[string]Service)
		s.startedServices = make(chan struct{}, 1)

		if s.Log == nil {
			s.Log = func(msg interface{}) {
				log.Printf("%s: %v", s.Name, msg)
			}
		}
		if s.FailureDecay == 0 {
			s.FailureDecay = 30
		}
		if s.FailureThreshold == 0 {
			s.FailureThreshold = 5
		}
		if s.Backoff == 0 {
			s.Backoff = 15 * time.Second
		}
	})
}

// Add inserts into the Supervisor tree a new service. If the Supervisor is
// already started, it will start it automatically.
func (s *Supervisor) Add(service Service) {
	s.prepare()

	name := fmt.Sprintf("%s", service)

	s.servicesMu.Lock()
	s.backoffMu.Lock()
	s.backoff[name] = &backoff{}
	s.services[name] = service
	s.backoffMu.Unlock()
	s.servicesMu.Unlock()

	select {
	case s.addedService <- struct{}{}:
	default:
	}
}

// Remove stops the service in the Supervisor tree and remove from it.
func (s *Supervisor) Remove(name string) {
	s.prepare()

	s.servicesMu.Lock()
	defer s.servicesMu.Unlock()
	if _, ok := s.services[name]; !ok {
		return
	}

	s.cancelationsMu.Lock()
	defer s.cancelationsMu.Unlock()
	if c, ok := s.cancelations[name]; ok {
		delete(s.cancelations, name)
		c()
	}
}

// Serve starts the Supervisor tree. It can be started only once at a time. If
// stopped (canceled), it can be restarted. In case of concurrent calls, it will
// hang until the current call is completed.
func (s *Supervisor) Serve(ctx context.Context) {
	s.prepare()

	s.runningMu.Lock()
	s.serve(ctx)
	s.runningMu.Unlock()
}

// Services return a list of services
func (s *Supervisor) Services() map[string]Service {
	svclist := make(map[string]Service)
	s.servicesMu.Lock()
	for k, v := range s.services {
		svclist[k] = v
	}
	s.servicesMu.Unlock()
	return svclist
}

// Cancelations return a list of services names and their cancellation calls
func (s *Supervisor) Cancelations() map[string]context.CancelFunc {
	svclist := make(map[string]context.CancelFunc)
	s.cancelationsMu.Lock()
	for k, v := range s.cancelations {
		svclist[k] = v
	}
	s.cancelationsMu.Unlock()
	return svclist
}

func (s *Supervisor) startServices(ctx context.Context) {
	s.startAllServices(ctx)
	select {
	case s.startedServices <- struct{}{}:
	default:
	}
}

func (s *Supervisor) serve(ctx context.Context) {
	select {
	case <-s.addedService:
	default:
	}
	s.startServices(ctx)

	go func(ctx context.Context) {
		for {
			select {
			case <-s.addedService:
				s.startServices(ctx)

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

	s.runningServices.Wait()

	s.cancelationsMu.Lock()
	s.cancelations = make(map[string]context.CancelFunc)
	s.cancelationsMu.Unlock()
	return
}

func (s *Supervisor) startAllServices(ctx context.Context) {
	s.servicesMu.Lock()
	defer s.servicesMu.Unlock()
	s.cancelationsMu.Lock()
	defer s.cancelationsMu.Unlock()

	var wg sync.WaitGroup
	for name, svc := range s.services {
		if _, ok := s.cancelations[name]; ok {
			continue
		}

		wg.Add(1)
		go func(name string, svc Service) {
			s.runningServices.Add(1)
			wg.Done()
			for {
				retry := func() (retry bool) {
					defer func() {
						if r := recover(); r != nil {
							s.Log(fmt.Sprintf("trapped panic: %v", r))
							retry = true
						}
					}()

					c, cancel := context.WithCancel(ctx)
					s.cancelationsMu.Lock()
					s.cancelations[name] = cancel
					s.cancelationsMu.Unlock()
					svc.Serve(c)

					select {
					case <-ctx.Done():
						return false
					default:
						return true
					}
				}()
				if retry {
					s.Log(fmt.Sprintf("restarting %s", name))
					s.backoffMu.Lock()
					b := s.backoff[name]
					s.backoffMu.Unlock()
					b.wait(s.FailureDecay, s.FailureThreshold, s.Backoff, func(msg interface{}) {
						s.Log(fmt.Sprintf("backoff %s: %v", name, msg))
					})
					continue
				}
				break
			}
			s.runningServices.Done()
		}(name, svc)
	}

	wg.Wait()
}

type backoff struct {
	lastfail time.Time
	failures float64
}

func (b *backoff) wait(failureDecay float64, threshold float64, backoffDur time.Duration, log func(str interface{})) {
	if b.lastfail.IsZero() {
		b.lastfail = time.Now()
		b.failures = 1.0
		return
	}

	b.failures++
	intervals := time.Since(b.lastfail).Seconds() / failureDecay
	b.failures = b.failures*math.Pow(.5, intervals) + 1

	if b.failures > threshold {
		log(backoffDur)
		time.Sleep(backoffDur)
	}
}
<|endoftext|>