<commit_before>package main

import (
	"net"
	"net/http"
	"net/url"
	"os/exec"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/mikesimons/earl"
	"github.com/mikesimons/pacyak/pacsandbox"
	"github.com/mikesimons/pacyak/proxyfactory"
	"github.com/mikesimons/readly"
)

const DIRECT_SANDBOX = 0
const PROXY_SANDBOX = 1

// PacYakOpts holds runtime config options for PacYakApplication
type PacYakOpts struct {
	PingCheckHost string
	PacFile       string
	ListenAddr    string
	PacProxy      string
	LogLevelStr   string
	LogLevel      log.Level
}

// pacInterpreter is a simple interface we use to provide a dummy implementation of pacsandbox for directPac
type pacInterpreter interface {
	ProxyFor(string) (string, error)
	Reset() // HACK
}

// directPac is a dummy implementation of pacsandbox to avoid invoking JS to return a constant static string ("DIRECT")
type directPac struct{}

func (p *directPac) ProxyFor(s string) (string, error) { return "DIRECT", nil }
func (p *directPac) Reset()                            {}

// PacYakApplication holds all application state
type PacYakApplication struct {
	opts         *PacYakOpts
	pacFile      *earl.URL
	sandboxIndex int
	sandboxes    []pacInterpreter
	factory      *proxyfactory.ProxyFactory
	listenAddr   string
	interfaceMap map[string]string
	Reader       *readly.Reader
}

// Run is the entry point for pacyak. It will initialize pacyak and start listening.
func Run(opts *PacYakOpts) {

	log.SetLevel(opts.LogLevel)
	reader := readly.New()

	// We need to explicitly set HTTP client to prevent it trying to use ENV vars for proxy
	// pacyak listen addr is expected to be set as HTTP_PROXY / HTTPS_PROXY but it isn't started yet!
	// This level of control also means a lib like hashicorp/go-getter is not suitable :(
	reader.Client = &http.Client{
		Transport: &http.Transport{
			Proxy: func(req *http.Request) (*url.URL, error) {
				if opts.PacProxy != "" {
					return url.Parse(opts.PacProxy)
				}

				return nil, nil
			},
			DialContext: (&net.Dialer{
				Timeout:   5 * time.Second,
				KeepAlive: 5 * time.Second,
			}).DialContext,
			IdleConnTimeout: 5 * time.Second,
		},
	}

	app := &PacYakApplication{
		opts:         opts,
		pacFile:      earl.Parse(opts.PacFile),
		factory:      proxyfactory.New(),
		sandboxIndex: DIRECT_SANDBOX,
		sandboxes:    []pacInterpreter{&directPac{}, &directPac{}},
		listenAddr:   opts.ListenAddr,
		Reader:       reader,
	}

	go app.monitorPingAvailability()
	go app.monitorNetworkInterfaces()

	// FIXME - graceful handler; server 502 on error and keep going
	log.Fatal(http.ListenAndServe(app.opts.ListenAddr, app))
}

// ServeHTTP handles directing the request to the correct proxy
func (app *PacYakApplication) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.WithFields(log.Fields{
		"method": r.Method,
		"url":    r.URL.String(),
	}).Debug("Processing HTTP request")

	// FIXME: We want to be able to serve some stats / tools from here
	//if !r.URL.IsAbs() {
	//	fmt.Fprintf(w, `function FindProxyForURL(url, host) { return "PROXY %s"; }`, app.listenAddr)
	//	return
	//}

	pacResponse, err := app.activeSandbox().ProxyFor(r.URL.String())

	if err != nil {
		log.WithFields(log.Fields{"response": pacResponse, "sandbox_error": err, "url": r.URL.String()}).Error("Sandbox error!")
	} else {
		log.WithFields(log.Fields{"response": pacResponse}).Debug("PAC result")
	}

	proxy := app.factory.FromPacResponse(pacResponse)

	proxy.ServeHTTP(w, r)
}

// switchToDirect switches the pac sandbox to the dummy "DIRECT" implementation
// We do this when our ping check fails (indicating a proxy may no longer be required)
func (app *PacYakApplication) switchToDirect() {
	if app.sandboxIndex != DIRECT_SANDBOX {
		log.Info("PAC availability check failed; switching to direct")
		app.setSandbox(DIRECT_SANDBOX)
	}
}

// switchToPac switches the pac sandbox to the JS implementation (using the PAC file specified on the CLI)
// We do this when our ping check passes (indicating we're in an env that requires a proxy)
func (app *PacYakApplication) switchToPac() {
	if app.sandboxIndex == DIRECT_SANDBOX {
		pac, err := app.Reader.Read(app.pacFile.Input)
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Error("PAC availability check passed but was unable to fetch PAC")
		} else {
			log.Info("PAC availability check passed; switching from direct")
			app.sandboxes[PROXY_SANDBOX] = pacsandbox.New(pac)
			app.setSandbox(PROXY_SANDBOX)
		}
	}
}

// handlePacAvailability updates the status of the ping check and switches sandbox if required
// This may be called from multiple go routines so we wrap it in a mutex to avoid racing
// Tried this with channels once but CPU usage blew up! Probably PEBKAC
func (app *PacYakApplication) handlePacAvailability() {
	available := exec.Command("ping", "-w", "1", app.opts.PingCheckHost).Run() == nil
	log.WithFields(log.Fields{"available": available}).Info("PAC availability check")

	if !available {
		app.switchToDirect()
	} else {
		app.switchToPac()
	}
}

// monitorPingAvailability is a wrapper for handlePacAvailability invoking it every 30 seconds
func (app *PacYakApplication) monitorPingAvailability() {
	app.handlePacAvailability()
	for _ = range time.Tick(30 * time.Second) {
		app.handlePacAvailability()
	}
}

// checkNetworkInterfaces will trigger a ping check if network interfaces have changed since last check
func (app *PacYakApplication) checkNetworkInterfaces() {
	interfaceMap := makeInterfaceMap()
	lastInterfaceMap := app.interfaceMap

	defer func() {
		app.interfaceMap = interfaceMap
	}()

	newInterfaces := interfaceMapKeys(interfaceMap)
	oldInterfaces := interfaceMapKeys(lastInterfaceMap)

	if interfaceListChanged(newInterfaces, oldInterfaces) {
		log.WithFields(log.Fields{"old": oldInterfaces, "new": newInterfaces}).Debug("Network interface list has changed")
		app.handlePacAvailability()
		return
	}

	for key, val := range interfaceMap {
		if lastInterfaceMap[key] != val {
			log.WithFields(log.Fields{"interface": key, "old": lastInterfaceMap[key], "new": val}).Debug("Network interface configuration has changed")
			app.handlePacAvailability()
			return
		}
	}

	log.Debug("No network changes detected")
}

// monitorNetworkInterfaces is a wrapper for checkNetworkInterfaces invoking it every 5 seconds
func (app *PacYakApplication) monitorNetworkInterfaces() {
	app.interfaceMap = makeInterfaceMap()
	for _ = range time.Tick(5 * time.Second) {
		app.checkNetworkInterfaces()
	}
}

func (app *PacYakApplication) activeSandbox() pacInterpreter {
	return app.sandboxes[app.sandboxIndex]
}

func (app *PacYakApplication) setSandbox(index int) {
	app.sandboxes[index].Reset()
	app.sandboxIndex = index
}
<commit_msg>Simple retry on PAC availability fail<commit_after>package main

import (
	"net"
	"net/http"
	"net/url"
	"os/exec"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/mikesimons/earl"
	"github.com/mikesimons/pacyak/pacsandbox"
	"github.com/mikesimons/pacyak/proxyfactory"
	"github.com/mikesimons/readly"
)

const DIRECT_SANDBOX = 0
const PROXY_SANDBOX = 1

// PacYakOpts holds runtime config options for PacYakApplication
type PacYakOpts struct {
	PingCheckHost string
	PacFile       string
	ListenAddr    string
	PacProxy      string
	LogLevelStr   string
	LogLevel      log.Level
}

// pacInterpreter is a simple interface we use to provide a dummy implementation of pacsandbox for directPac
type pacInterpreter interface {
	ProxyFor(string) (string, error)
	Reset() // HACK
}

// directPac is a dummy implementation of pacsandbox to avoid invoking JS to return a constant static string ("DIRECT")
type directPac struct{}

func (p *directPac) ProxyFor(s string) (string, error) { return "DIRECT", nil }
func (p *directPac) Reset()                            {}

// PacYakApplication holds all application state
type PacYakApplication struct {
	opts         *PacYakOpts
	pacFile      *earl.URL
	sandboxIndex int
	sandboxes    []pacInterpreter
	factory      *proxyfactory.ProxyFactory
	listenAddr   string
	interfaceMap map[string]string
	Reader       *readly.Reader
}

// Run is the entry point for pacyak. It will initialize pacyak and start listening.
func Run(opts *PacYakOpts) {

	log.SetLevel(opts.LogLevel)
	reader := readly.New()

	// We need to explicitly set HTTP client to prevent it trying to use ENV vars for proxy
	// pacyak listen addr is expected to be set as HTTP_PROXY / HTTPS_PROXY but it isn't started yet!
	// This level of control also means a lib like hashicorp/go-getter is not suitable :(
	reader.Client = &http.Client{
		Transport: &http.Transport{
			Proxy: func(req *http.Request) (*url.URL, error) {
				if opts.PacProxy != "" {
					return url.Parse(opts.PacProxy)
				}

				return nil, nil
			},
			DialContext: (&net.Dialer{
				Timeout:   5 * time.Second,
				KeepAlive: 5 * time.Second,
			}).DialContext,
			IdleConnTimeout: 5 * time.Second,
		},
	}

	app := &PacYakApplication{
		opts:         opts,
		pacFile:      earl.Parse(opts.PacFile),
		factory:      proxyfactory.New(),
		sandboxIndex: DIRECT_SANDBOX,
		sandboxes:    []pacInterpreter{&directPac{}, &directPac{}},
		listenAddr:   opts.ListenAddr,
		Reader:       reader,
	}

	go app.monitorPingAvailability()
	go app.monitorNetworkInterfaces()

	// FIXME - graceful handler; server 502 on error and keep going
	log.Fatal(http.ListenAndServe(app.opts.ListenAddr, app))
}

// ServeHTTP handles directing the request to the correct proxy
func (app *PacYakApplication) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	log.WithFields(log.Fields{
		"method": r.Method,
		"url":    r.URL.String(),
	}).Debug("Processing HTTP request")

	// FIXME: We want to be able to serve some stats / tools from here
	//if !r.URL.IsAbs() {
	//	fmt.Fprintf(w, `function FindProxyForURL(url, host) { return "PROXY %s"; }`, app.listenAddr)
	//	return
	//}

	pacResponse, err := app.activeSandbox().ProxyFor(r.URL.String())

	if err != nil {
		log.WithFields(log.Fields{"response": pacResponse, "sandbox_error": err, "url": r.URL.String()}).Error("Sandbox error!")
	} else {
		log.WithFields(log.Fields{"response": pacResponse}).Debug("PAC result")
	}

	proxy := app.factory.FromPacResponse(pacResponse)

	proxy.ServeHTTP(w, r)
}

// switchToDirect switches the pac sandbox to the dummy "DIRECT" implementation
// We do this when our ping check fails (indicating a proxy may no longer be required)
func (app *PacYakApplication) switchToDirect() {
	if app.sandboxIndex != DIRECT_SANDBOX {
		log.Info("PAC availability check failed; switching to direct")
		app.setSandbox(DIRECT_SANDBOX)
	}
}

// switchToPac switches the pac sandbox to the JS implementation (using the PAC file specified on the CLI)
// We do this when our ping check passes (indicating we're in an env that requires a proxy)
func (app *PacYakApplication) switchToPac() {
	if app.sandboxIndex == DIRECT_SANDBOX {
		pac, err := app.Reader.Read(app.pacFile.Input)
		if err != nil {
			log.WithFields(log.Fields{"error": err}).Error("PAC availability check passed but was unable to fetch PAC")
		} else {
			log.Info("PAC availability check passed; switching from direct")
			app.sandboxes[PROXY_SANDBOX] = pacsandbox.New(pac)
			app.setSandbox(PROXY_SANDBOX)
		}
	}
}

// handlePacAvailability updates the status of the ping check and switches sandbox if required
// This may be called from multiple go routines so we wrap it in a mutex to avoid racing
// Tried this with channels once but CPU usage blew up! Probably PEBKAC
func (app *PacYakApplication) handlePacAvailability() {
	available := false
	retries := 0
	for ; retries < 2; retries++ {
		available = exec.Command("ping", "-w", "1", app.opts.PingCheckHost).Run() == nil
		log.WithFields(log.Fields{"available": available}).Info("PAC availability check")

		if !available {
			time.Sleep(time.Second * 5)
		} else {
			break
		}
	}

	if !available {
		app.switchToDirect()
	} else {
		app.switchToPac()
	}
}

// monitorPingAvailability is a wrapper for handlePacAvailability invoking it every 30 seconds
func (app *PacYakApplication) monitorPingAvailability() {
	app.handlePacAvailability()
	for _ = range time.Tick(30 * time.Second) {
		app.handlePacAvailability()
	}
}

// checkNetworkInterfaces will trigger a ping check if network interfaces have changed since last check
func (app *PacYakApplication) checkNetworkInterfaces() {
	interfaceMap := makeInterfaceMap()
	lastInterfaceMap := app.interfaceMap

	defer func() {
		app.interfaceMap = interfaceMap
	}()

	newInterfaces := interfaceMapKeys(interfaceMap)
	oldInterfaces := interfaceMapKeys(lastInterfaceMap)

	if interfaceListChanged(newInterfaces, oldInterfaces) {
		log.WithFields(log.Fields{"old": oldInterfaces, "new": newInterfaces}).Debug("Network interface list has changed")
		app.handlePacAvailability()
		return
	}

	for key, val := range interfaceMap {
		if lastInterfaceMap[key] != val {
			log.WithFields(log.Fields{"interface": key, "old": lastInterfaceMap[key], "new": val}).Debug("Network interface configuration has changed")
			app.handlePacAvailability()
			return
		}
	}

	log.Debug("No network changes detected")
}

// monitorNetworkInterfaces is a wrapper for checkNetworkInterfaces invoking it every 5 seconds
func (app *PacYakApplication) monitorNetworkInterfaces() {
	app.interfaceMap = makeInterfaceMap()
	for _ = range time.Tick(5 * time.Second) {
		app.checkNetworkInterfaces()
	}
}

func (app *PacYakApplication) activeSandbox() pacInterpreter {
	return app.sandboxes[app.sandboxIndex]
}

func (app *PacYakApplication) setSandbox(index int) {
	app.sandboxes[index].Reset()
	app.sandboxIndex = index
}
<|endoftext|>


<commit_before>// +build darwin

package process

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"unsafe"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/internal/common"
	"github.com/shirou/gopsutil/net"
)

// copied from sys/sysctl.h
const (
	CTLKern          = 1  // "high kernel": proc, limits
	KernProc         = 14 // struct: process entries
	KernProcPID      = 1  // by process id
	KernProcProc     = 8  // only return procs
	KernProcAll      = 0  // everything
	KernProcPathname = 12 // path to executable
)

const (
	ClockTicks = 100 // C.sysconf(C._SC_CLK_TCK)
)

type _Ctype_struct___0 struct {
	Pad uint64
}

// MemoryInfoExStat is different between OSes
type MemoryInfoExStat struct {
}

type MemoryMapsStat struct {
}

func Pids() ([]int32, error) {
	var ret []int32

	pids, err := callPs("pid", 0, false)
	if err != nil {
		return ret, err
	}

	for _, pid := range pids {
		v, err := strconv.Atoi(pid[0])
		if err != nil {
			return ret, err
		}
		ret = append(ret, int32(v))
	}

	return ret, nil
}

func (p *Process) Ppid() (int32, error) {
	r, err := callPs("ppid", p.Pid, false)
	v, err := strconv.Atoi(r[0][0])
	if err != nil {
		return 0, err
	}

	return int32(v), err
}
func (p *Process) Name() (string, error) {
	k, err := p.getKProc()
	if err != nil {
		return "", err
	}

	return common.IntToString(k.Proc.P_comm[:]), nil
}
func (p *Process) Exe() (string, error) {
	return "", common.ErrNotImplementedError
}

// Cmdline returns the command line arguments of the process as a string with
// each argument separated by 0x20 ascii character.
func (p *Process) Cmdline() (string, error) {
	r, err := callPs("command", p.Pid, false)
	if err != nil {
		return "", err
	}
	return strings.Join(r[0], " "), err
}

// CmdlineSlice returns the command line arguments of the process as a slice with each
// element being an argument. Because of current deficiencies in the way that the command
// line arguments are found, single arguments that have spaces in the will actually be
// reported as two separate items. In order to do something better CGO would be needed
// to use the native darwin functions.
func (p *Process) CmdlineSlice() ([]string, error) {
	r, err := callPs("command", p.Pid, false)
	if err != nil {
		return nil, err
	}
	return r[0], err
}
func (p *Process) CreateTime() (int64, error) {
	return 0, common.ErrNotImplementedError
}
func (p *Process) Cwd() (string, error) {
	return "", common.ErrNotImplementedError
}
func (p *Process) Parent() (*Process, error) {
	rr, err := common.CallLsof(invoke, p.Pid, "-FR")
	if err != nil {
		return nil, err
	}
	for _, r := range rr {
		if strings.HasPrefix(r, "p") { // skip if process
			continue
		}
		l := string(r)
		v, err := strconv.Atoi(strings.Replace(l, "R", "", 1))
		if err != nil {
			return nil, err
		}
		return NewProcess(int32(v))
	}
	return nil, fmt.Errorf("could not find parent line")
}
func (p *Process) Status() (string, error) {
	r, err := callPs("state", p.Pid, false)
	if err != nil {
		return "", err
	}

	return r[0][0], err
}
func (p *Process) Uids() ([]int32, error) {
	k, err := p.getKProc()
	if err != nil {
		return nil, err
	}

	// See: http://unix.superglobalmegacorp.com/Net2/newsrc/sys/ucred.h.html
	userEffectiveUID := int32(k.Eproc.Ucred.UID)

	return []int32{userEffectiveUID}, nil
}
func (p *Process) Gids() ([]int32, error) {
	k, err := p.getKProc()
	if err != nil {
		return nil, err
	}

	gids := make([]int32, 0, 3)
	gids = append(gids, int32(k.Eproc.Pcred.P_rgid), int32(k.Eproc.Ucred.Ngroups), int32(k.Eproc.Pcred.P_svgid))

	return gids, nil
}
func (p *Process) Terminal() (string, error) {
	return "", common.ErrNotImplementedError
	/*
		k, err := p.getKProc()
		if err != nil {
			return "", err
		}

		ttyNr := uint64(k.Eproc.Tdev)
		termmap, err := getTerminalMap()
		if err != nil {
			return "", err
		}

		return termmap[ttyNr], nil
	*/
}
func (p *Process) Nice() (int32, error) {
	k, err := p.getKProc()
	if err != nil {
		return 0, err
	}
	return int32(k.Proc.P_nice), nil
}
func (p *Process) IOnice() (int32, error) {
	return 0, common.ErrNotImplementedError
}
func (p *Process) Rlimit() ([]RlimitStat, error) {
	var rlimit []RlimitStat
	return rlimit, common.ErrNotImplementedError
}
func (p *Process) IOCounters() (*IOCountersStat, error) {
	return nil, common.ErrNotImplementedError
}
func (p *Process) NumCtxSwitches() (*NumCtxSwitchesStat, error) {
	return nil, common.ErrNotImplementedError
}
func (p *Process) NumFDs() (int32, error) {
	return 0, common.ErrNotImplementedError
}
func (p *Process) NumThreads() (int32, error) {
	r, err := callPs("utime,stime", p.Pid, true)
	if err != nil {
		return 0, err
	}
	return int32(len(r)), nil
}
func (p *Process) Threads() (map[string]string, error) {
	ret := make(map[string]string, 0)
	return ret, common.ErrNotImplementedError
}

func convertCPUTimes(s string) (ret float64, err error) {
	var t int
	var _tmp string
	if strings.Contains(s, ":") {
		_t := strings.Split(s, ":")
		hour, err := strconv.Atoi(_t[0])
		if err != nil {
			return ret, err
		}
		t += hour * 60 * 100
		_tmp = _t[1]
	} else {
		_tmp = s
	}

	_t := strings.Split(_tmp, ".")
	if err != nil {
		return ret, err
	}
	h, err := strconv.Atoi(_t[0])
	t += h * 100
	h, err = strconv.Atoi(_t[1])
	t += h
	return float64(t) / ClockTicks, nil
}
func (p *Process) Times() (*cpu.TimesStat, error) {
	r, err := callPs("utime,stime", p.Pid, false)

	if err != nil {
		return nil, err
	}

	utime, err := convertCPUTimes(r[0][0])
	if err != nil {
		return nil, err
	}
	stime, err := convertCPUTimes(r[0][1])
	if err != nil {
		return nil, err
	}

	ret := &cpu.TimesStat{
		CPU:    "cpu",
		User:   utime,
		System: stime,
	}
	return ret, nil
}
func (p *Process) CPUAffinity() ([]int32, error) {
	return nil, common.ErrNotImplementedError
}
func (p *Process) MemoryInfo() (*MemoryInfoStat, error) {
	r, err := callPs("rss,vsize,pagein", p.Pid, false)
	if err != nil {
		return nil, err
	}
	rss, err := strconv.Atoi(r[0][0])
	if err != nil {
		return nil, err
	}
	vms, err := strconv.Atoi(r[0][1])
	if err != nil {
		return nil, err
	}
	pagein, err := strconv.Atoi(r[0][2])
	if err != nil {
		return nil, err
	}

	ret := &MemoryInfoStat{
		RSS:  uint64(rss) * 1024,
		VMS:  uint64(vms) * 1024,
		Swap: uint64(pagein),
	}

	return ret, nil
}
func (p *Process) MemoryInfoEx() (*MemoryInfoExStat, error) {
	return nil, common.ErrNotImplementedError
}

func (p *Process) Children() ([]*Process, error) {
	pids, err := common.CallPgrep(invoke, p.Pid)
	if err != nil {
		return nil, err
	}
	ret := make([]*Process, 0, len(pids))
	for _, pid := range pids {
		np, err := NewProcess(pid)
		if err != nil {
			return nil, err
		}
		ret = append(ret, np)
	}
	return ret, nil
}

func (p *Process) OpenFiles() ([]OpenFilesStat, error) {
	return nil, common.ErrNotImplementedError
}

func (p *Process) Connections() ([]net.ConnectionStat, error) {
	return net.ConnectionsPid("all", p.Pid)
}

func (p *Process) NetIOCounters(pernic bool) ([]net.IOCountersStat, error) {
	return nil, common.ErrNotImplementedError
}

func (p *Process) IsRunning() (bool, error) {
	return true, common.ErrNotImplementedError
}
func (p *Process) MemoryMaps(grouped bool) (*[]MemoryMapsStat, error) {
	var ret []MemoryMapsStat
	return &ret, common.ErrNotImplementedError
}

func processes() ([]Process, error) {
	results := make([]Process, 0, 50)

	mib := []int32{CTLKern, KernProc, KernProcAll, 0}
	buf, length, err := common.CallSyscall(mib)
	if err != nil {
		return results, err
	}

	// get kinfo_proc size
	k := KinfoProc{}
	procinfoLen := int(unsafe.Sizeof(k))
	count := int(length / uint64(procinfoLen))
	/*
		fmt.Println(length, procinfoLen, count)
		b := buf[0*procinfoLen : 0*procinfoLen+procinfoLen]
		fmt.Println(b)
		kk, err := parseKinfoProc(b)
		fmt.Printf("%#v", kk)
	*/

	// parse buf to procs
	for i := 0; i < count; i++ {
		b := buf[i*procinfoLen : i*procinfoLen+procinfoLen]
		k, err := parseKinfoProc(b)
		if err != nil {
			continue
		}
		p, err := NewProcess(int32(k.Proc.P_pid))
		if err != nil {
			continue
		}
		results = append(results, *p)
	}

	return results, nil
}

func parseKinfoProc(buf []byte) (KinfoProc, error) {
	var k KinfoProc
	br := bytes.NewReader(buf)

	err := common.Read(br, binary.LittleEndian, &k)
	if err != nil {
		return k, err
	}

	return k, nil
}

// Returns a proc as defined here:
// http://unix.superglobalmegacorp.com/Net2/newsrc/sys/kinfo_proc.h.html
func (p *Process) getKProc() (*KinfoProc, error) {
	mib := []int32{CTLKern, KernProc, KernProcPID, p.Pid}
	procK := KinfoProc{}
	length := uint64(unsafe.Sizeof(procK))
	buf := make([]byte, length)
	_, _, syserr := syscall.Syscall6(
		syscall.SYS___SYSCTL,
		uintptr(unsafe.Pointer(&mib[0])),
		uintptr(len(mib)),
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&length)),
		0,
		0)
	if syserr != 0 {
		return nil, syserr
	}
	k, err := parseKinfoProc(buf)
	if err != nil {
		return nil, err
	}

	return &k, nil
}

func NewProcess(pid int32) (*Process, error) {
	p := &Process{Pid: pid}

	return p, nil
}

// call ps command.
// Return value deletes Header line(you must not input wrong arg).
// And splited by Space. Caller have responsibility to manage.
// If passed arg pid is 0, get information from all process.
func callPs(arg string, pid int32, threadOption bool) ([][]string, error) {
	bin, err := exec.LookPath("ps")
	if err != nil {
		return [][]string{}, err
	}

	var cmd []string
	if pid == 0 { // will get from all processes.
		cmd = []string{"-ax", "-o", arg}
	} else if threadOption {
		cmd = []string{"-x", "-o", arg, "-M", "-p", strconv.Itoa(int(pid))}
	} else {
		cmd = []string{"-x", "-o", arg, "-p", strconv.Itoa(int(pid))}
	}
	out, err := invoke.Command(bin, cmd...)
	if err != nil {
		return [][]string{}, err
	}
	lines := strings.Split(string(out), "\n")

	var ret [][]string
	for _, l := range lines[1:] {
		var lr []string
		for _, r := range strings.Split(l, " ") {
			if r == "" {
				continue
			}
			lr = append(lr, strings.TrimSpace(r))
		}
		if len(lr) != 0 {
			ret = append(ret, lr)
		}
	}

	return ret, nil
}
<commit_msg>Implement CreateTime for darwin. Parses elapsed time and converts this to epoch time.<commit_after>// +build darwin

package process

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unsafe"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/internal/common"
	"github.com/shirou/gopsutil/net"
)

// copied from sys/sysctl.h
const (
	CTLKern          = 1  // "high kernel": proc, limits
	KernProc         = 14 // struct: process entries
	KernProcPID      = 1  // by process id
	KernProcProc     = 8  // only return procs
	KernProcAll      = 0  // everything
	KernProcPathname = 12 // path to executable
)

const (
	ClockTicks = 100 // C.sysconf(C._SC_CLK_TCK)
)

type _Ctype_struct___0 struct {
	Pad uint64
}

// MemoryInfoExStat is different between OSes
type MemoryInfoExStat struct {
}

type MemoryMapsStat struct {
}

func Pids() ([]int32, error) {
	var ret []int32

	pids, err := callPs("pid", 0, false)
	if err != nil {
		return ret, err
	}

	for _, pid := range pids {
		v, err := strconv.Atoi(pid[0])
		if err != nil {
			return ret, err
		}
		ret = append(ret, int32(v))
	}

	return ret, nil
}

func (p *Process) Ppid() (int32, error) {
	r, err := callPs("ppid", p.Pid, false)
	v, err := strconv.Atoi(r[0][0])
	if err != nil {
		return 0, err
	}

	return int32(v), err
}
func (p *Process) Name() (string, error) {
	k, err := p.getKProc()
	if err != nil {
		return "", err
	}

	return common.IntToString(k.Proc.P_comm[:]), nil
}
func (p *Process) Exe() (string, error) {
	return "", common.ErrNotImplementedError
}

// Cmdline returns the command line arguments of the process as a string with
// each argument separated by 0x20 ascii character.
func (p *Process) Cmdline() (string, error) {
	r, err := callPs("command", p.Pid, false)
	if err != nil {
		return "", err
	}
	return strings.Join(r[0], " "), err
}

// CmdlineSlice returns the command line arguments of the process as a slice with each
// element being an argument. Because of current deficiencies in the way that the command
// line arguments are found, single arguments that have spaces in the will actually be
// reported as two separate items. In order to do something better CGO would be needed
// to use the native darwin functions.
func (p *Process) CmdlineSlice() ([]string, error) {
	r, err := callPs("command", p.Pid, false)
	if err != nil {
		return nil, err
	}
	return r[0], err
}
func (p *Process) CreateTime() (int64, error) {
	r, err := callPs("etime", p.Pid, false)
	if err != nil {
		return 0, err
	}

	elapsedSegments := strings.Split(r[0][0], ":")
	var elapsedDurations []time.Duration
	for i := len(elapsedSegments) - 1; i >= 0; i-- {
		p, err := strconv.ParseInt(elapsedSegments[i], 10, 0)
		if err != nil {
			return 0, err
		}
		elapsedDurations = append(elapsedDurations, time.Duration(p))
	}

	var elapsed time.Duration = time.Duration(elapsedDurations[0]) * time.Second
	if len(elapsedDurations) > 1 {
		elapsed += time.Duration(elapsedDurations[1]) * time.Minute
	}
	if len(elapsedDurations) > 2 {
		elapsed += time.Duration(elapsedDurations[2]) * time.Hour
	}
	if len(elapsedDurations) > 3 {
		elapsed += time.Duration(elapsedDurations[3]) * time.Hour * 24
	}

	start := time.Now().Add(-elapsed)
	return start.Unix() * 1000, nil
}
func (p *Process) Cwd() (string, error) {
	return "", common.ErrNotImplementedError
}
func (p *Process) Parent() (*Process, error) {
	rr, err := common.CallLsof(invoke, p.Pid, "-FR")
	if err != nil {
		return nil, err
	}
	for _, r := range rr {
		if strings.HasPrefix(r, "p") { // skip if process
			continue
		}
		l := string(r)
		v, err := strconv.Atoi(strings.Replace(l, "R", "", 1))
		if err != nil {
			return nil, err
		}
		return NewProcess(int32(v))
	}
	return nil, fmt.Errorf("could not find parent line")
}
func (p *Process) Status() (string, error) {
	r, err := callPs("state", p.Pid, false)
	if err != nil {
		return "", err
	}

	return r[0][0], err
}
func (p *Process) Uids() ([]int32, error) {
	k, err := p.getKProc()
	if err != nil {
		return nil, err
	}

	// See: http://unix.superglobalmegacorp.com/Net2/newsrc/sys/ucred.h.html
	userEffectiveUID := int32(k.Eproc.Ucred.UID)

	return []int32{userEffectiveUID}, nil
}
func (p *Process) Gids() ([]int32, error) {
	k, err := p.getKProc()
	if err != nil {
		return nil, err
	}

	gids := make([]int32, 0, 3)
	gids = append(gids, int32(k.Eproc.Pcred.P_rgid), int32(k.Eproc.Ucred.Ngroups), int32(k.Eproc.Pcred.P_svgid))

	return gids, nil
}
func (p *Process) Terminal() (string, error) {
	return "", common.ErrNotImplementedError
	/*
		k, err := p.getKProc()
		if err != nil {
			return "", err
		}

		ttyNr := uint64(k.Eproc.Tdev)
		termmap, err := getTerminalMap()
		if err != nil {
			return "", err
		}

		return termmap[ttyNr], nil
	*/
}
func (p *Process) Nice() (int32, error) {
	k, err := p.getKProc()
	if err != nil {
		return 0, err
	}
	return int32(k.Proc.P_nice), nil
}
func (p *Process) IOnice() (int32, error) {
	return 0, common.ErrNotImplementedError
}
func (p *Process) Rlimit() ([]RlimitStat, error) {
	var rlimit []RlimitStat
	return rlimit, common.ErrNotImplementedError
}
func (p *Process) IOCounters() (*IOCountersStat, error) {
	return nil, common.ErrNotImplementedError
}
func (p *Process) NumCtxSwitches() (*NumCtxSwitchesStat, error) {
	return nil, common.ErrNotImplementedError
}
func (p *Process) NumFDs() (int32, error) {
	return 0, common.ErrNotImplementedError
}
func (p *Process) NumThreads() (int32, error) {
	r, err := callPs("utime,stime", p.Pid, true)
	if err != nil {
		return 0, err
	}
	return int32(len(r)), nil
}
func (p *Process) Threads() (map[string]string, error) {
	ret := make(map[string]string, 0)
	return ret, common.ErrNotImplementedError
}

func convertCPUTimes(s string) (ret float64, err error) {
	var t int
	var _tmp string
	if strings.Contains(s, ":") {
		_t := strings.Split(s, ":")
		hour, err := strconv.Atoi(_t[0])
		if err != nil {
			return ret, err
		}
		t += hour * 60 * 100
		_tmp = _t[1]
	} else {
		_tmp = s
	}

	_t := strings.Split(_tmp, ".")
	if err != nil {
		return ret, err
	}
	h, err := strconv.Atoi(_t[0])
	t += h * 100
	h, err = strconv.Atoi(_t[1])
	t += h
	return float64(t) / ClockTicks, nil
}
func (p *Process) Times() (*cpu.TimesStat, error) {
	r, err := callPs("utime,stime", p.Pid, false)

	if err != nil {
		return nil, err
	}

	utime, err := convertCPUTimes(r[0][0])
	if err != nil {
		return nil, err
	}
	stime, err := convertCPUTimes(r[0][1])
	if err != nil {
		return nil, err
	}

	ret := &cpu.TimesStat{
		CPU:    "cpu",
		User:   utime,
		System: stime,
	}
	return ret, nil
}
func (p *Process) CPUAffinity() ([]int32, error) {
	return nil, common.ErrNotImplementedError
}
func (p *Process) MemoryInfo() (*MemoryInfoStat, error) {
	r, err := callPs("rss,vsize,pagein", p.Pid, false)
	if err != nil {
		return nil, err
	}
	rss, err := strconv.Atoi(r[0][0])
	if err != nil {
		return nil, err
	}
	vms, err := strconv.Atoi(r[0][1])
	if err != nil {
		return nil, err
	}
	pagein, err := strconv.Atoi(r[0][2])
	if err != nil {
		return nil, err
	}

	ret := &MemoryInfoStat{
		RSS:  uint64(rss) * 1024,
		VMS:  uint64(vms) * 1024,
		Swap: uint64(pagein),
	}

	return ret, nil
}
func (p *Process) MemoryInfoEx() (*MemoryInfoExStat, error) {
	return nil, common.ErrNotImplementedError
}

func (p *Process) Children() ([]*Process, error) {
	pids, err := common.CallPgrep(invoke, p.Pid)
	if err != nil {
		return nil, err
	}
	ret := make([]*Process, 0, len(pids))
	for _, pid := range pids {
		np, err := NewProcess(pid)
		if err != nil {
			return nil, err
		}
		ret = append(ret, np)
	}
	return ret, nil
}

func (p *Process) OpenFiles() ([]OpenFilesStat, error) {
	return nil, common.ErrNotImplementedError
}

func (p *Process) Connections() ([]net.ConnectionStat, error) {
	return net.ConnectionsPid("all", p.Pid)
}

func (p *Process) NetIOCounters(pernic bool) ([]net.IOCountersStat, error) {
	return nil, common.ErrNotImplementedError
}

func (p *Process) IsRunning() (bool, error) {
	return true, common.ErrNotImplementedError
}
func (p *Process) MemoryMaps(grouped bool) (*[]MemoryMapsStat, error) {
	var ret []MemoryMapsStat
	return &ret, common.ErrNotImplementedError
}

func processes() ([]Process, error) {
	results := make([]Process, 0, 50)

	mib := []int32{CTLKern, KernProc, KernProcAll, 0}
	buf, length, err := common.CallSyscall(mib)
	if err != nil {
		return results, err
	}

	// get kinfo_proc size
	k := KinfoProc{}
	procinfoLen := int(unsafe.Sizeof(k))
	count := int(length / uint64(procinfoLen))
	/*
		fmt.Println(length, procinfoLen, count)
		b := buf[0*procinfoLen : 0*procinfoLen+procinfoLen]
		fmt.Println(b)
		kk, err := parseKinfoProc(b)
		fmt.Printf("%#v", kk)
	*/

	// parse buf to procs
	for i := 0; i < count; i++ {
		b := buf[i*procinfoLen : i*procinfoLen+procinfoLen]
		k, err := parseKinfoProc(b)
		if err != nil {
			continue
		}
		p, err := NewProcess(int32(k.Proc.P_pid))
		if err != nil {
			continue
		}
		results = append(results, *p)
	}

	return results, nil
}

func parseKinfoProc(buf []byte) (KinfoProc, error) {
	var k KinfoProc
	br := bytes.NewReader(buf)

	err := common.Read(br, binary.LittleEndian, &k)
	if err != nil {
		return k, err
	}

	return k, nil
}

// Returns a proc as defined here:
// http://unix.superglobalmegacorp.com/Net2/newsrc/sys/kinfo_proc.h.html
func (p *Process) getKProc() (*KinfoProc, error) {
	mib := []int32{CTLKern, KernProc, KernProcPID, p.Pid}
	procK := KinfoProc{}
	length := uint64(unsafe.Sizeof(procK))
	buf := make([]byte, length)
	_, _, syserr := syscall.Syscall6(
		syscall.SYS___SYSCTL,
		uintptr(unsafe.Pointer(&mib[0])),
		uintptr(len(mib)),
		uintptr(unsafe.Pointer(&buf[0])),
		uintptr(unsafe.Pointer(&length)),
		0,
		0)
	if syserr != 0 {
		return nil, syserr
	}
	k, err := parseKinfoProc(buf)
	if err != nil {
		return nil, err
	}

	return &k, nil
}

func NewProcess(pid int32) (*Process, error) {
	p := &Process{Pid: pid}

	return p, nil
}

// call ps command.
// Return value deletes Header line(you must not input wrong arg).
// And splited by Space. Caller have responsibility to manage.
// If passed arg pid is 0, get information from all process.
func callPs(arg string, pid int32, threadOption bool) ([][]string, error) {
	bin, err := exec.LookPath("ps")
	if err != nil {
		return [][]string{}, err
	}

	var cmd []string
	if pid == 0 { // will get from all processes.
		cmd = []string{"-ax", "-o", arg}
	} else if threadOption {
		cmd = []string{"-x", "-o", arg, "-M", "-p", strconv.Itoa(int(pid))}
	} else {
		cmd = []string{"-x", "-o", arg, "-p", strconv.Itoa(int(pid))}
	}
	out, err := invoke.Command(bin, cmd...)
	if err != nil {
		return [][]string{}, err
	}
	lines := strings.Split(string(out), "\n")

	var ret [][]string
	for _, l := range lines[1:] {
		var lr []string
		for _, r := range strings.Split(l, " ") {
			if r == "" {
				continue
			}
			lr = append(lr, strings.TrimSpace(r))
		}
		if len(lr) != 0 {
			ret = append(ret, lr)
		}
	}

	return ret, nil
}
<|endoftext|>


<commit_before>package store

import (
	"encoding/json"
	"fmt"
	"github.com/Cepave/open-falcon-backend/common/model"
	"github.com/Cepave/open-falcon-backend/modules/judge/g"
	"log"
)

func Judge(L *SafeLinkedList, firstItem *model.JudgeItem, now int64) {
	CheckStrategy(L, firstItem, now)
	CheckExpression(L, firstItem, now)
}

func CheckStrategy(L *SafeLinkedList, firstItem *model.JudgeItem, now int64) {
	key := fmt.Sprintf("%s/%s", firstItem.Endpoint, firstItem.Metric)
	strategyMap := g.StrategyMap.Get()
	strategies, exists := strategyMap[key]
	if !exists {
		return
	}

	for _, s := range strategies {
		// 因为key仅仅是endpoint和metric，所以得到的strategies并不一定是与当前judgeItem相关的
		// 比如lg-dinp-docker01.bj配置了两个proc.num的策略，一个name=docker，一个name=agent
		// 所以此处要排除掉一部分
		related := true
		for tagKey, tagVal := range s.Tags {
			if myVal, exists := firstItem.Tags[tagKey]; !exists || myVal != tagVal {
				related = false
				break
			}
		}

		if !related {
			continue
		}

		judgeItemWithStrategy(L, s, firstItem, now)
	}
}

func judgeItemWithStrategy(L *SafeLinkedList, strategy model.Strategy, firstItem *model.JudgeItem, now int64) {
	fn, err := ParseFuncFromString(strategy.Func, strategy.Operator, strategy.RightValue)
	if err != nil {
		log.Printf("[ERROR] parse func %s fail: %v. strategy id: %d", strategy.Func, err, strategy.Id)
		return
	}

	historyData, leftValue, isTriggered, isEnough := fn.Compute(L)
	if !isEnough {
		return
	}

	event := &model.Event{
		Id:         fmt.Sprintf("s_%d_%s", strategy.Id, firstItem.PrimaryKey()),
		Strategy:   &strategy,
		Endpoint:   firstItem.Endpoint,
		LeftValue:  leftValue,
		EventTime:  firstItem.Timestamp,
		PushedTags: firstItem.Tags,
	}

	sendEventIfNeed(historyData, isTriggered, now, event, strategy.MaxStep)
}

func sendEvent(event *model.Event) {
	// update last event
	g.LastEvents.Set(event.Id, event)

	bs, err := json.Marshal(event)
	if err != nil {
		log.Printf("json marshal event %v fail: %v", event, err)
		return
	}

	// send to redis
	redisKey := fmt.Sprintf(g.Config().Alarm.QueuePattern, event.Priority())
	rc := g.RedisConnPool.Get()
	defer rc.Close()
	rc.Do("LPUSH", redisKey, string(bs))
}

func CheckExpression(L *SafeLinkedList, firstItem *model.JudgeItem, now int64) {
	keys := buildKeysFromMetricAndTags(firstItem)
	if len(keys) == 0 {
		return
	}

	// expression可能会被多次重复处理，用此数据结构保证只被处理一次
	handledExpression := make(map[int]struct{})

	expressionMap := g.ExpressionMap.Get()
	for _, key := range keys {
		expressions, exists := expressionMap[key]
		if !exists {
			continue
		}

		related := filterRelatedExpressions(expressions, firstItem)
		for _, exp := range related {
			if _, ok := handledExpression[exp.Id]; ok {
				continue
			}
			handledExpression[exp.Id] = struct{}{}
			judgeItemWithExpression(L, exp, firstItem, now)
		}
	}
}

func buildKeysFromMetricAndTags(item *model.JudgeItem) (keys []string) {
	for k, v := range item.Tags {
		keys = append(keys, fmt.Sprintf("%s/%s=%s", item.Metric, k, v))
	}
	keys = append(keys, fmt.Sprintf("%s/endpoint=%s", item.Metric, item.Endpoint))
	return
}

func filterRelatedExpressions(expressions []*model.Expression, firstItem *model.JudgeItem) []*model.Expression {
	size := len(expressions)
	if size == 0 {
		return []*model.Expression{}
	}

	exps := make([]*model.Expression, 0, size)

	for _, exp := range expressions {

		related := true

		itemTagsCopy := firstItem.Tags
		// 注意：exp.Tags 中可能会有一个endpoint=xxx的tag
		if _, ok := exp.Tags["endpoint"]; ok {
			itemTagsCopy = copyItemTags(firstItem)
		}

		for tagKey, tagVal := range exp.Tags {
			if myVal, exists := itemTagsCopy[tagKey]; !exists || myVal != tagVal {
				related = false
				break
			}
		}

		if !related {
			continue
		}

		exps = append(exps, exp)
	}

	return exps
}

func copyItemTags(item *model.JudgeItem) map[string]string {
	ret := make(map[string]string)
	ret["endpoint"] = item.Endpoint
	if item.Tags != nil && len(item.Tags) > 0 {
		for k, v := range item.Tags {
			ret[k] = v
		}
	}
	return ret
}

func judgeItemWithExpression(L *SafeLinkedList, expression *model.Expression, firstItem *model.JudgeItem, now int64) {
	fn, err := ParseFuncFromString(expression.Func, expression.Operator, expression.RightValue)
	if err != nil {
		log.Printf("[ERROR] parse func %s fail: %v. expression id: %d", expression.Func, err, expression.Id)
		return
	}

	historyData, leftValue, isTriggered, isEnough := fn.Compute(L)
	if !isEnough {
		return
	}

	event := &model.Event{
		Id:         fmt.Sprintf("e_%d_%s", expression.Id, firstItem.PrimaryKey()),
		Expression: expression,
		Endpoint:   firstItem.Endpoint,
		LeftValue:  leftValue,
		EventTime:  firstItem.Timestamp,
		PushedTags: firstItem.Tags,
	}

	sendEventIfNeed(historyData, isTriggered, now, event, expression.MaxStep)

}

func sendEventIfNeed(historyData []*model.HistoryData, isTriggered bool, now int64, event *model.Event, maxStep int) {
	lastEvent, exists := g.LastEvents.Get(event.Id)
	if isTriggered {
		event.Status = "PROBLEM"
		if !exists || lastEvent.Status[0] == 'O' {
			// 本次触发了阈值，之前又没报过警，得产生一个报警Event
			event.CurrentStep = 1

			// 但是有些用户把最大报警次数配置成了0，相当于屏蔽了，要检查一下
			if maxStep == 0 {
				return
			}

			sendEvent(event)
			return
		}

		// 逻辑走到这里，说明之前Event是PROBLEM状态
		if lastEvent.CurrentStep >= maxStep {
			// 报警次数已经足够多，到达了最多报警次数了，不再报警
			return
		}

		if historyData[len(historyData)-1].Timestamp <= lastEvent.EventTime {
			// 产生过报警的点，就不能再使用来判断了，否则容易出现一分钟报一次的情况
			// 只需要拿最后一个historyData来做判断即可，因为它的时间最老
			return
		}

		if now-lastEvent.EventTime < g.Config().Alarm.MinInterval {
			// 报警不能太频繁，两次报警之间至少要间隔MinInterval秒，否则就不能报警
			return
		}

		event.CurrentStep = lastEvent.CurrentStep + 1
		sendEvent(event)
	} else {
		// 如果LastEvent是Problem，报OK，否则啥都不做
		if exists && lastEvent.Status[0] == 'P' {
			event.Status = "OK"
			event.CurrentStep = 1
			sendEvent(event)
		}
	}
}
<commit_msg>add mechanism to re-set current_step , when strategy is changed <commit_after>package store

import (
	"encoding/json"
	"fmt"
	"log"

	"github.com/Cepave/open-falcon-backend/common/model"
	"github.com/Cepave/open-falcon-backend/modules/judge/g"
)

func Judge(L *SafeLinkedList, firstItem *model.JudgeItem, now int64) {
	CheckStrategy(L, firstItem, now)
	CheckExpression(L, firstItem, now)
}

func CheckStrategy(L *SafeLinkedList, firstItem *model.JudgeItem, now int64) {
	key := fmt.Sprintf("%s/%s", firstItem.Endpoint, firstItem.Metric)
	strategyMap := g.StrategyMap.Get()
	strategies, exists := strategyMap[key]
	if !exists {
		return
	}

	for _, s := range strategies {
		// 因为key仅仅是endpoint和metric，所以得到的strategies并不一定是与当前judgeItem相关的
		// 比如lg-dinp-docker01.bj配置了两个proc.num的策略，一个name=docker，一个name=agent
		// 所以此处要排除掉一部分
		related := true
		for tagKey, tagVal := range s.Tags {
			if myVal, exists := firstItem.Tags[tagKey]; !exists || myVal != tagVal {
				related = false
				break
			}
		}

		if !related {
			continue
		}

		judgeItemWithStrategy(L, s, firstItem, now)
	}
}

func judgeItemWithStrategy(L *SafeLinkedList, strategy model.Strategy, firstItem *model.JudgeItem, now int64) {
	fn, err := ParseFuncFromString(strategy.Func, strategy.Operator, strategy.RightValue)
	if err != nil {
		log.Printf("[ERROR] parse func %s fail: %v. strategy id: %d", strategy.Func, err, strategy.Id)
		return
	}

	historyData, leftValue, isTriggered, isEnough := fn.Compute(L)
	if !isEnough {
		return
	}

	event := &model.Event{
		Id:         fmt.Sprintf("s_%d_%s", strategy.Id, firstItem.PrimaryKey()),
		Strategy:   &strategy,
		Endpoint:   firstItem.Endpoint,
		LeftValue:  leftValue,
		EventTime:  firstItem.Timestamp,
		PushedTags: firstItem.Tags,
	}

	sendEventIfNeed(historyData, isTriggered, now, event, strategy.MaxStep)
}

func sendEvent(event *model.Event) {
	// update last event
	g.LastEvents.Set(event.Id, event)

	bs, err := json.Marshal(event)
	if err != nil {
		log.Printf("json marshal event %v fail: %v", event, err)
		return
	}

	// send to redis
	redisKey := fmt.Sprintf(g.Config().Alarm.QueuePattern, event.Priority())
	rc := g.RedisConnPool.Get()
	defer rc.Close()
	rc.Do("LPUSH", redisKey, string(bs))
}

func CheckExpression(L *SafeLinkedList, firstItem *model.JudgeItem, now int64) {
	keys := buildKeysFromMetricAndTags(firstItem)
	if len(keys) == 0 {
		return
	}

	// expression可能会被多次重复处理，用此数据结构保证只被处理一次
	handledExpression := make(map[int]struct{})

	expressionMap := g.ExpressionMap.Get()
	for _, key := range keys {
		expressions, exists := expressionMap[key]
		if !exists {
			continue
		}

		related := filterRelatedExpressions(expressions, firstItem)
		for _, exp := range related {
			if _, ok := handledExpression[exp.Id]; ok {
				continue
			}
			handledExpression[exp.Id] = struct{}{}
			judgeItemWithExpression(L, exp, firstItem, now)
		}
	}
}

func buildKeysFromMetricAndTags(item *model.JudgeItem) (keys []string) {
	for k, v := range item.Tags {
		keys = append(keys, fmt.Sprintf("%s/%s=%s", item.Metric, k, v))
	}
	keys = append(keys, fmt.Sprintf("%s/endpoint=%s", item.Metric, item.Endpoint))
	return
}

func filterRelatedExpressions(expressions []*model.Expression, firstItem *model.JudgeItem) []*model.Expression {
	size := len(expressions)
	if size == 0 {
		return []*model.Expression{}
	}

	exps := make([]*model.Expression, 0, size)

	for _, exp := range expressions {

		related := true

		itemTagsCopy := firstItem.Tags
		// 注意：exp.Tags 中可能会有一个endpoint=xxx的tag
		if _, ok := exp.Tags["endpoint"]; ok {
			itemTagsCopy = copyItemTags(firstItem)
		}

		for tagKey, tagVal := range exp.Tags {
			if myVal, exists := itemTagsCopy[tagKey]; !exists || myVal != tagVal {
				related = false
				break
			}
		}

		if !related {
			continue
		}

		exps = append(exps, exp)
	}

	return exps
}

func copyItemTags(item *model.JudgeItem) map[string]string {
	ret := make(map[string]string)
	ret["endpoint"] = item.Endpoint
	if item.Tags != nil && len(item.Tags) > 0 {
		for k, v := range item.Tags {
			ret[k] = v
		}
	}
	return ret
}

func judgeItemWithExpression(L *SafeLinkedList, expression *model.Expression, firstItem *model.JudgeItem, now int64) {
	fn, err := ParseFuncFromString(expression.Func, expression.Operator, expression.RightValue)
	if err != nil {
		log.Printf("[ERROR] parse func %s fail: %v. expression id: %d", expression.Func, err, expression.Id)
		return
	}

	historyData, leftValue, isTriggered, isEnough := fn.Compute(L)
	if !isEnough {
		return
	}

	event := &model.Event{
		Id:         fmt.Sprintf("e_%d_%s", expression.Id, firstItem.PrimaryKey()),
		Expression: expression,
		Endpoint:   firstItem.Endpoint,
		LeftValue:  leftValue,
		EventTime:  firstItem.Timestamp,
		PushedTags: firstItem.Tags,
	}

	sendEventIfNeed(historyData, isTriggered, now, event, expression.MaxStep)

}

func sendEventIfNeed(historyData []*model.HistoryData, isTriggered bool, now int64, event *model.Event, maxStep int) {
	lastEvent, exists := g.LastEvents.Get(event.Id)
	needSet := false
	if exists {
		//if strategy is changed will reset currentStep to "1"
		needSet = fmt.Sprintf("%s %v", lastEvent.Strategy.Operator, lastEvent.Strategy.RightValue) != fmt.Sprintf("%s %v", event.Strategy.Operator, event.Strategy.RightValue)
	}
	if isTriggered {
		event.Status = "PROBLEM"
		if !exists || lastEvent.Status[0] == 'O' || needSet {
			// 本次触发了阈值，之前又没报过警，得产生一个报警Event
			event.CurrentStep = 1

			// 但是有些用户把最大报警次数配置成了0，相当于屏蔽了，要检查一下
			if maxStep == 0 {
				return
			}

			sendEvent(event)
			return
		}

		// 逻辑走到这里，说明之前Event是PROBLEM状态
		if lastEvent.CurrentStep >= maxStep {
			// 报警次数已经足够多，到达了最多报警次数了，不再报警
			return
		}

		if historyData[len(historyData)-1].Timestamp <= lastEvent.EventTime {
			// 产生过报警的点，就不能再使用来判断了，否则容易出现一分钟报一次的情况
			// 只需要拿最后一个historyData来做判断即可，因为它的时间最老
			return
		}

		if now-lastEvent.EventTime < g.Config().Alarm.MinInterval {
			// 报警不能太频繁，两次报警之间至少要间隔MinInterval秒，否则就不能报警
			return
		}

		event.CurrentStep = lastEvent.CurrentStep + 1
		sendEvent(event)
	} else {
		// 如果LastEvent是Problem，报OK，否则啥都不做
		if exists && lastEvent.Status[0] == 'P' {
			event.Status = "OK"
			event.CurrentStep = 1
			sendEvent(event)
		}
	}
}
<|endoftext|>


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

package lfs

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/storage"
)

var (
	errHashMismatch = errors.New("Content hash does not match OID")
	errSizeMismatch = errors.New("Content size does not match")
)

// ErrRangeNotSatisfiable represents an error which request range is not satisfiable.
type ErrRangeNotSatisfiable struct {
	FromByte int64
}

func (err ErrRangeNotSatisfiable) Error() string {
	return fmt.Sprintf("Requested range %d is not satisfiable", err.FromByte)
}

// IsErrRangeNotSatisfiable returns true if the error is an ErrRangeNotSatisfiable
func IsErrRangeNotSatisfiable(err error) bool {
	_, ok := err.(ErrRangeNotSatisfiable)
	return ok
}

// ContentStore provides a simple file system based storage.
type ContentStore struct {
	storage.ObjectStorage
}

// Get takes a Meta object and retrieves the content from the store, returning
// it as an io.Reader. If fromByte > 0, the reader starts from that byte
func (s *ContentStore) Get(meta *models.LFSMetaObject, fromByte int64) (io.ReadCloser, error) {
	f, err := s.Open(meta.RelativePath())
	if err != nil {
		log.Error("Whilst trying to read LFS OID[%s]: Unable to open Error: %v", meta.Oid, err)
		return nil, err
	}
	if fromByte > 0 {
		if fromByte >= meta.Size {
			return nil, ErrRangeNotSatisfiable{
				FromByte: fromByte,
			}
		}
		_, err = f.Seek(fromByte, io.SeekStart)
		if err != nil {
			log.Error("Whilst trying to read LFS OID[%s]: Unable to seek to %d Error: %v", meta.Oid, fromByte, err)
		}
	}
	return f, err
}

// Put takes a Meta object and an io.Reader and writes the content to the store.
func (s *ContentStore) Put(meta *models.LFSMetaObject, r io.Reader) error {
	hash := sha256.New()
	rd := io.TeeReader(r, hash)
	p := meta.RelativePath()
	written, err := s.Save(p, rd)
	if err != nil {
		log.Error("Whilst putting LFS OID[%s]: Failed to copy to tmpPath: %s Error: %v", meta.Oid, p, err)
		return err
	}

	if written != meta.Size {
		if err := s.Delete(p); err != nil {
			log.Error("Cleaning the LFS OID[%s] failed: %v", meta.Oid, err)
		}
		return errSizeMismatch
	}

	shaStr := hex.EncodeToString(hash.Sum(nil))
	if shaStr != meta.Oid {
		if err := s.Delete(p); err != nil {
			log.Error("Cleaning the LFS OID[%s] failed: %v", meta.Oid, err)
		}
		return errHashMismatch
	}

	return nil
}

// Exists returns true if the object exists in the content store.
func (s *ContentStore) Exists(meta *models.LFSMetaObject) (bool, error) {
	_, err := s.ObjectStorage.Stat(meta.RelativePath())
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// Verify returns true if the object exists in the content store and size is correct.
func (s *ContentStore) Verify(meta *models.LFSMetaObject) (bool, error) {
	p := meta.RelativePath()
	fi, err := s.ObjectStorage.Stat(p)
	if os.IsNotExist(err) || (err == nil && fi.Size() != meta.Size) {
		return false, nil
	} else if err != nil {
		log.Error("Unable stat file: %s for LFS OID[%s] Error: %v", p, meta.Oid, err)
		return false, err
	}

	return true, nil
}
<commit_msg>Fix race in LFS ContentStore.Put(...) (#14895)<commit_after>// Copyright 2020 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package lfs

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"hash"
	"io"
	"os"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/storage"
)

var (
	errHashMismatch = errors.New("Content hash does not match OID")
	errSizeMismatch = errors.New("Content size does not match")
)

// ErrRangeNotSatisfiable represents an error which request range is not satisfiable.
type ErrRangeNotSatisfiable struct {
	FromByte int64
}

func (err ErrRangeNotSatisfiable) Error() string {
	return fmt.Sprintf("Requested range %d is not satisfiable", err.FromByte)
}

// IsErrRangeNotSatisfiable returns true if the error is an ErrRangeNotSatisfiable
func IsErrRangeNotSatisfiable(err error) bool {
	_, ok := err.(ErrRangeNotSatisfiable)
	return ok
}

// ContentStore provides a simple file system based storage.
type ContentStore struct {
	storage.ObjectStorage
}

// Get takes a Meta object and retrieves the content from the store, returning
// it as an io.Reader. If fromByte > 0, the reader starts from that byte
func (s *ContentStore) Get(meta *models.LFSMetaObject, fromByte int64) (io.ReadCloser, error) {
	f, err := s.Open(meta.RelativePath())
	if err != nil {
		log.Error("Whilst trying to read LFS OID[%s]: Unable to open Error: %v", meta.Oid, err)
		return nil, err
	}
	if fromByte > 0 {
		if fromByte >= meta.Size {
			return nil, ErrRangeNotSatisfiable{
				FromByte: fromByte,
			}
		}
		_, err = f.Seek(fromByte, io.SeekStart)
		if err != nil {
			log.Error("Whilst trying to read LFS OID[%s]: Unable to seek to %d Error: %v", meta.Oid, fromByte, err)
		}
	}
	return f, err
}

// Put takes a Meta object and an io.Reader and writes the content to the store.
func (s *ContentStore) Put(meta *models.LFSMetaObject, r io.Reader) error {
	p := meta.RelativePath()

	// Wrap the provided reader with an inline hashing and size checker
	wrappedRd := newHashingReader(meta.Size, meta.Oid, r)

	// now pass the wrapped reader to Save - if there is a size mismatch or hash mismatch then
	// the errors returned by the newHashingReader should percolate up to here
	written, err := s.Save(p, wrappedRd)
	if err != nil {
		log.Error("Whilst putting LFS OID[%s]: Failed to copy to tmpPath: %s Error: %v", meta.Oid, p, err)
		return err
	}

	// This shouldn't happen but it is sensible to test
	if written != meta.Size {
		if err := s.Delete(p); err != nil {
			log.Error("Cleaning the LFS OID[%s] failed: %v", meta.Oid, err)
		}
		return errSizeMismatch
	}

	return nil
}

// Exists returns true if the object exists in the content store.
func (s *ContentStore) Exists(meta *models.LFSMetaObject) (bool, error) {
	_, err := s.ObjectStorage.Stat(meta.RelativePath())
	if err != nil {
		if os.IsNotExist(err) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// Verify returns true if the object exists in the content store and size is correct.
func (s *ContentStore) Verify(meta *models.LFSMetaObject) (bool, error) {
	p := meta.RelativePath()
	fi, err := s.ObjectStorage.Stat(p)
	if os.IsNotExist(err) || (err == nil && fi.Size() != meta.Size) {
		return false, nil
	} else if err != nil {
		log.Error("Unable stat file: %s for LFS OID[%s] Error: %v", p, meta.Oid, err)
		return false, err
	}

	return true, nil
}

type hashingReader struct {
	internal     io.Reader
	currentSize  int64
	expectedSize int64
	hash         hash.Hash
	expectedHash string
}

func (r *hashingReader) Read(b []byte) (int, error) {
	n, err := r.internal.Read(b)

	if n > 0 {
		r.currentSize += int64(n)
		wn, werr := r.hash.Write(b[:n])
		if wn != n || werr != nil {
			return n, werr
		}
	}

	if err != nil && err == io.EOF {
		if r.currentSize != r.expectedSize {
			return n, errSizeMismatch
		}

		shaStr := hex.EncodeToString(r.hash.Sum(nil))
		if shaStr != r.expectedHash {
			return n, errHashMismatch
		}
	}

	return n, err
}

func newHashingReader(expectedSize int64, expectedHash string, reader io.Reader) *hashingReader {
	return &hashingReader{
		internal:     reader,
		expectedSize: expectedSize,
		expectedHash: expectedHash,
		hash:         sha256.New(),
	}
}
<|endoftext|>


<commit_before>package sourcegraph

import (
	"crypto/tls"
	"net"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/oauth"
)

type contextKey int

const (
	grpcEndpointKey contextKey = iota
	httpEndpointKey
	credentialsKey
	clientMetadataKey
)

// WithGRPCEndpoint returns a copy of parent whose clients (obtained
// using FromContext) communicate with the given gRPC API endpoint
// URL.
func WithGRPCEndpoint(parent context.Context, url *url.URL) context.Context {
	return context.WithValue(parent, grpcEndpointKey, url)
}

// GRPCEndpoint returns the context's gRPC endpoint URL that was
// previously configured using WithGRPCEndpoint.
func GRPCEndpoint(ctx context.Context) *url.URL {
	url, _ := ctx.Value(grpcEndpointKey).(*url.URL)
	if url == nil {
		panic("no gRPC API endpoint URL set in context")
	}
	return url
}

// Credentials authenticate gRPC requests made by an API client.
type Credentials interface {
	oauth2.TokenSource
}

// WithCredentials returns a copy of the parent context that uses cred
// as the credentials for future API clients constructed using this
// context (with NewClientFromContext). It replaces (shadows) any
// previously set credentials in the context.
//
// It can be used to add, e.g., trace/span ID metadata for request
// tracing.
func WithCredentials(parent context.Context, cred Credentials) context.Context {
	return context.WithValue(parent, credentialsKey, cred)
}

// CredentialsFromContext returns the credentials (if any) previously
// set in the context by WithCredentials.
func CredentialsFromContext(ctx context.Context) Credentials {
	cred, ok := ctx.Value(credentialsKey).(Credentials)
	if !ok {
		return nil
	}
	return cred
}

// WithClientMetadata returns a copy of the parent context that merges
// in the specified metadata to future API clients constructed using
// this context (with NewClientFromContext). It replaces (shadows) any
// previously set metadata in the context.
func WithClientMetadata(parent context.Context, md map[string]string) context.Context {
	return context.WithValue(parent, clientMetadataKey, md)
}

// clientMetadataFromContext returns the metadata (if any) previously
// set in the context by WithClientMetadata.
func clientMetadataFromContext(ctx context.Context) map[string]string {
	cred, ok := ctx.Value(clientMetadataKey).(map[string]string)
	if !ok {
		return nil
	}
	return cred
}

var maxDialTimeout = 10 * time.Second

// NewClientFromContext returns a Sourcegraph API client that
// communicates with the Sourcegraph gRPC endpoint in ctx (i.e.,
// GRPCEndpoint(ctx)).
var NewClientFromContext = func(ctx context.Context) *Client {
	opts := []grpc.DialOption{
		grpc.WithCodec(GRPCCodec),
	}

	grpcEndpoint := GRPCEndpoint(ctx)
	if grpcEndpoint.Scheme == "https" {
		creds := credentials.NewClientTLSFromCert(nil, "")
		if host, _, _ := net.SplitHostPort(grpcEndpoint.Host); host == "localhost" {
			creds = credentials.NewTLS(&tls.Config{InsecureSkipVerify: true})
		}
		opts = append(opts, grpc.WithTransportCredentials(creds))
	} else {
		opts = append(opts, grpc.WithInsecure())
	}

	// Use contextCredentials instead of directly using the cred
	// so that we can use different credentials for the same
	// connection (in the pool).
	opts = append(opts, grpc.WithPerRPCCredentials(contextCredentials{}))

	// Dial timeout is the lesser of the ctx deadline or
	// maxDialTimeout.
	var timeout time.Duration
	if d, ok := ctx.Deadline(); ok && time.Now().Add(maxDialTimeout).After(d) {
		timeout = d.Sub(time.Now())
	} else {
		timeout = maxDialTimeout
	}
	opts = append(opts, grpc.WithTimeout(timeout))

	conn, err := pooledGRPCDial(grpcEndpoint.Host, opts...)
	if err != nil {
		panic(err)
	}
	c := NewClient(conn)
	return c
}

// RemovePooledGRPCConn removes the pooled grpc.ClientConnection to the gRPC endpoint
// in the context. The result of calling this function  is that the pooled connection for
// this endpoint will be reset, so the subsequent call to NewClientFromContext() would have
// to dial a new gRPC connection to this endpoint.
var RemovePooledGRPCConn = func(ctx context.Context) {
	grpcEndpoint := GRPCEndpoint(ctx)
	removeConnFromPool(grpcEndpoint.Host)
}

// contextCredentials implements the credentials.Credentials interface.
type contextCredentials struct{}

// GetRequestMetadata implements the credentials.Credentials interface. As per the
// interface definition, it may be called by multiple goroutines concurrently.
func (contextCredentials) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
	m := clientMetadataFromContext(ctx)

	if cred := CredentialsFromContext(ctx); cred != nil {
		credMD, err := (oauth.TokenSource{TokenSource: cred}).GetRequestMetadata(ctx)
		if err != nil {
			return nil, err
		}

		if m == nil {
			m = credMD
		} else {
			// Copy the map to avoid a data race writing to the map returned by
			// clientMetaDataFromContext.
			cpy := make(map[string]string)
			for k, v := range m {
				cpy[k] = v
			}
			for k, v := range credMD {
				cpy[k] = v
			}
			m = cpy
		}
	}
	return m, nil
}

// RequireTransportSecurity implements the credentials.Credentials interface.
func (contextCredentials) RequireTransportSecurity() bool {
	return false
}
<commit_msg>race: Synchronise access to NewClientFromContext<commit_after>package sourcegraph

import (
	"crypto/tls"
	"net"
	"net/url"
	"sync"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/oauth2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/oauth"
)

type contextKey int

const (
	grpcEndpointKey contextKey = iota
	httpEndpointKey
	credentialsKey
	clientMetadataKey
)

// WithGRPCEndpoint returns a copy of parent whose clients (obtained
// using FromContext) communicate with the given gRPC API endpoint
// URL.
func WithGRPCEndpoint(parent context.Context, url *url.URL) context.Context {
	return context.WithValue(parent, grpcEndpointKey, url)
}

// GRPCEndpoint returns the context's gRPC endpoint URL that was
// previously configured using WithGRPCEndpoint.
func GRPCEndpoint(ctx context.Context) *url.URL {
	url, _ := ctx.Value(grpcEndpointKey).(*url.URL)
	if url == nil {
		panic("no gRPC API endpoint URL set in context")
	}
	return url
}

// Credentials authenticate gRPC requests made by an API client.
type Credentials interface {
	oauth2.TokenSource
}

// WithCredentials returns a copy of the parent context that uses cred
// as the credentials for future API clients constructed using this
// context (with NewClientFromContext). It replaces (shadows) any
// previously set credentials in the context.
//
// It can be used to add, e.g., trace/span ID metadata for request
// tracing.
func WithCredentials(parent context.Context, cred Credentials) context.Context {
	return context.WithValue(parent, credentialsKey, cred)
}

// CredentialsFromContext returns the credentials (if any) previously
// set in the context by WithCredentials.
func CredentialsFromContext(ctx context.Context) Credentials {
	cred, ok := ctx.Value(credentialsKey).(Credentials)
	if !ok {
		return nil
	}
	return cred
}

// WithClientMetadata returns a copy of the parent context that merges
// in the specified metadata to future API clients constructed using
// this context (with NewClientFromContext). It replaces (shadows) any
// previously set metadata in the context.
func WithClientMetadata(parent context.Context, md map[string]string) context.Context {
	return context.WithValue(parent, clientMetadataKey, md)
}

// clientMetadataFromContext returns the metadata (if any) previously
// set in the context by WithClientMetadata.
func clientMetadataFromContext(ctx context.Context) map[string]string {
	cred, ok := ctx.Value(clientMetadataKey).(map[string]string)
	if !ok {
		return nil
	}
	return cred
}

// NewClientFromContext returns a Sourcegraph API client that
// communicates with the Sourcegraph gRPC endpoint in ctx (i.e.,
// GRPCEndpoint(ctx)).
func NewClientFromContext(ctx context.Context) *Client {
	newClientFromContextMu.RLock()
	f := newClientFromContext
	newClientFromContextMu.RUnlock()
	return f(ctx)
}

// MockNewClientFromContext allows a test to mock out the return value of
// NewClientFromContext. Note that this is modifying global state, so if your
// tests run in parallel you may get unexpected results
func MockNewClientFromContext(f func(ctx context.Context) *Client) {
	newClientFromContextMu.Lock()
	newClientFromContext = f
	newClientFromContextMu.Unlock()
}

// RestoreNewClientFromContext removes the mock and returns the correct
// implementation
func RestoreNewClientFromContext() {
	newClientFromContextMu.Lock()
	newClientFromContext = realNewClientFromContext
	newClientFromContextMu.Unlock()
}

var (
	maxDialTimeout         = 10 * time.Second
	newClientFromContextMu sync.RWMutex
	newClientFromContext   = realNewClientFromContext
)
var realNewClientFromContext = func(ctx context.Context) *Client {
	opts := []grpc.DialOption{
		grpc.WithCodec(GRPCCodec),
	}

	grpcEndpoint := GRPCEndpoint(ctx)
	if grpcEndpoint.Scheme == "https" {
		creds := credentials.NewClientTLSFromCert(nil, "")
		if host, _, _ := net.SplitHostPort(grpcEndpoint.Host); host == "localhost" {
			creds = credentials.NewTLS(&tls.Config{InsecureSkipVerify: true})
		}
		opts = append(opts, grpc.WithTransportCredentials(creds))
	} else {
		opts = append(opts, grpc.WithInsecure())
	}

	// Use contextCredentials instead of directly using the cred
	// so that we can use different credentials for the same
	// connection (in the pool).
	opts = append(opts, grpc.WithPerRPCCredentials(contextCredentials{}))

	// Dial timeout is the lesser of the ctx deadline or
	// maxDialTimeout.
	var timeout time.Duration
	if d, ok := ctx.Deadline(); ok && time.Now().Add(maxDialTimeout).After(d) {
		timeout = d.Sub(time.Now())
	} else {
		timeout = maxDialTimeout
	}
	opts = append(opts, grpc.WithTimeout(timeout))

	conn, err := pooledGRPCDial(grpcEndpoint.Host, opts...)
	if err != nil {
		panic(err)
	}
	c := NewClient(conn)
	return c
}

// RemovePooledGRPCConn removes the pooled grpc.ClientConnection to the gRPC endpoint
// in the context. The result of calling this function  is that the pooled connection for
// this endpoint will be reset, so the subsequent call to NewClientFromContext() would have
// to dial a new gRPC connection to this endpoint.
var RemovePooledGRPCConn = func(ctx context.Context) {
	grpcEndpoint := GRPCEndpoint(ctx)
	removeConnFromPool(grpcEndpoint.Host)
}

// contextCredentials implements the credentials.Credentials interface.
type contextCredentials struct{}

// GetRequestMetadata implements the credentials.Credentials interface. As per the
// interface definition, it may be called by multiple goroutines concurrently.
func (contextCredentials) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
	m := clientMetadataFromContext(ctx)

	if cred := CredentialsFromContext(ctx); cred != nil {
		credMD, err := (oauth.TokenSource{TokenSource: cred}).GetRequestMetadata(ctx)
		if err != nil {
			return nil, err
		}

		if m == nil {
			m = credMD
		} else {
			// Copy the map to avoid a data race writing to the map returned by
			// clientMetaDataFromContext.
			cpy := make(map[string]string)
			for k, v := range m {
				cpy[k] = v
			}
			for k, v := range credMD {
				cpy[k] = v
			}
			m = cpy
		}
	}
	return m, nil
}

// RequireTransportSecurity implements the credentials.Credentials interface.
func (contextCredentials) RequireTransportSecurity() bool {
	return false
}
<|endoftext|>


<commit_before>package amqp

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/brokers/errs"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/common"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/pkg/errors"
	"github.com/streadway/amqp"
)

type AMQPConnection struct {
	queueName    string
	connection   *amqp.Connection
	channel      *amqp.Channel
	queue        amqp.Queue
	confirmation <-chan amqp.Confirmation
	errorchan    <-chan *amqp.Error
	cleanup      chan struct{}
}

// Broker represents an AMQP broker
type Broker struct {
	common.Broker
	common.AMQPConnector
	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal

	connections      map[string]*AMQPConnection
	connectionsMutex sync.RWMutex
}

// New creates new Broker instance
func New(cnf *config.Config) iface.Broker {
	return &Broker{Broker: common.NewBroker(cnf), AMQPConnector: common.AMQPConnector{}, connections: make(map[string]*AMQPConnection)}
}

// StartConsuming enters a loop and waits for incoming messages
func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)

	queueName := taskProcessor.CustomQueue()
	if queueName == "" {
		queueName = b.GetConfig().DefaultQueue
	}

	conn, channel, queue, _, amqpCloseChan, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().MultipleBrokerSeparator,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		b.GetConfig().AMQP.BindingKey,   // queue binding key
		nil,                             // exchange declare args
		amqp.Table(b.GetConfig().AMQP.QueueDeclareArgs), // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		b.GetRetryFunc()(b.GetRetryStopChan())
		return b.GetRetry(), err
	}
	defer b.Close(channel, conn)

	if err = channel.Qos(
		b.GetConfig().AMQP.PrefetchCount,
		0,     // prefetch size
		false, // global
	); err != nil {
		return b.GetRetry(), fmt.Errorf("Channel qos error: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return b.GetRetry(), fmt.Errorf("Queue consume error: %s", err)
	}

	log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := b.consume(deliveries, concurrency, taskProcessor, amqpCloseChan); err != nil {
		return b.GetRetry(), err
	}

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	return b.GetRetry(), nil
}

// StopConsuming quits the loop
func (b *Broker) StopConsuming() {
	b.Broker.StopConsuming()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()
}

// GetOrOpenConnection will return a connection on a particular queue name. Open connections
// are saved to avoid having to reopen connection for multiple queues
func (b *Broker) GetOrOpenConnection(queueName string, queueBindingKey string, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs amqp.Table) (*AMQPConnection, error) {
	var err error

	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	conn, ok := b.connections[queueName]
	if !ok {
		conn = &AMQPConnection{
			queueName: queueName,
			cleanup:   make(chan struct{}),
		}
		conn.connection, conn.channel, conn.queue, conn.confirmation, conn.errorchan, err = b.Connect(
			b.GetConfig().Broker,
			b.GetConfig().MultipleBrokerSeparator,
			b.GetConfig().TLSConfig,
			b.GetConfig().AMQP.Exchange,     // exchange name
			b.GetConfig().AMQP.ExchangeType, // exchange type
			queueName,                       // queue name
			true,                            // queue durable
			false,                           // queue delete when unused
			queueBindingKey,                 // queue binding key
			exchangeDeclareArgs,             // exchange declare args
			queueDeclareArgs,                // queue declare args
			queueBindingArgs,                // queue binding args
		)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to connect to queue %s", queueName)
		}

		// Reconnect to the channel if it disconnects/errors out
		go func() {
			select {
			case err = <-conn.errorchan:
				log.INFO.Printf("Error occurred on queue: %s. Reconnecting", queueName)
				b.connectionsMutex.Lock()
				delete(b.connections, queueName)
				b.connectionsMutex.Unlock()
				_, err := b.GetOrOpenConnection(queueName, queueBindingKey, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs)
				if err != nil {
					log.ERROR.Printf("Failed to reopen queue: %s.", queueName)
				}
			case <-conn.cleanup:
				return
			}
			return
		}()
		b.connections[queueName] = conn
	}
	return conn, nil
}

func (b *Broker) CloseConnections() error {
	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	for key, conn := range b.connections {
		if err := b.Close(conn.channel, conn.connection); err != nil {
			log.ERROR.Print("Failed to close channel")
			return nil
		}
		close(conn.cleanup)
		delete(b.connections, key)
	}
	return nil
}

// Publish places a new message on the default queue
func (b *Broker) Publish(ctx context.Context, signature *tasks.Signature) error {
	// Adjust routing key (this decides which queue the message will be published to)
	b.AdjustRoutingKey(signature)

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// delay the task
	if signature.ETA != nil {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			delayMs := int64(signature.ETA.Sub(now) / time.Millisecond)

			return b.delay(signature, delayMs)
		}
	}

	queue := b.GetConfig().DefaultQueue
	bindingKey := b.GetConfig().AMQP.BindingKey // queue binding key
	if b.isDirectExchange() {
		queue = signature.RoutingKey
		bindingKey = signature.RoutingKey
	}

	connection, err := b.GetOrOpenConnection(
		queue,
		bindingKey, // queue binding key
		nil,        // exchange declare args
		amqp.Table(b.GetConfig().AMQP.QueueDeclareArgs), // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return errors.Wrapf(err, "Failed to get a connection for queue %s", queue)
	}

	channel := connection.channel
	confirmsChan := connection.confirmation

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange name
		signature.RoutingKey,        // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         msg,
			Priority:     signature.Priority,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return errors.Wrap(err, "Failed to publish task")
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// consume takes delivered messages from the channel and manages a worker pool
// to process tasks concurrently
func (b *Broker) consume(deliveries <-chan amqp.Delivery, concurrency int, taskProcessor iface.TaskProcessor, amqpCloseChan <-chan *amqp.Error) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		for i := 0; i < concurrency; i++ {
			pool <- struct{}{}
		}
	}()

	// make channel with a capacity makes it become a buffered channel so that a worker which wants to
	// push an error to `errorsChan` doesn't need to be blocked while the for-loop is blocked waiting
	// a worker, that is, it avoids a possible deadlock
	errorsChan := make(chan error, 1)

	for {
		select {
		case amqpErr := <-amqpCloseChan:
			return amqpErr
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			if concurrency > 0 {
				// get worker from pool (blocks until one is available)
				<-pool
			}

			b.processingWG.Add(1)

			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				if err := b.consumeOne(d, taskProcessor); err != nil {
					errorsChan <- err
				}

				b.processingWG.Done()

				if concurrency > 0 {
					// give worker back to pool
					pool <- struct{}{}
				}
			}()
		case <-b.GetStopChan():
			return nil
		}
	}
}

// consumeOne processes a single message using TaskProcessor
func (b *Broker) consumeOne(delivery amqp.Delivery, taskProcessor iface.TaskProcessor) error {
	if len(delivery.Body) == 0 {
		delivery.Nack(true, false)                     // multiple, requeue
		return errors.New("Received an empty message") // RabbitMQ down?
	}

	var multiple, requeue = false, false

	// Unmarshal message body into signature struct
	signature := new(tasks.Signature)
	decoder := json.NewDecoder(bytes.NewReader(delivery.Body))
	decoder.UseNumber()
	if err := decoder.Decode(signature); err != nil {
		delivery.Nack(multiple, requeue)
		return errs.NewErrCouldNotUnmarshaTaskSignature(delivery.Body, err)
	}

	// If the task is not registered, we nack it and requeue,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(signature.Name) {
		if !delivery.Redelivered {
			requeue = true
			log.INFO.Printf("Task not registered with this worker. Requeing message: %s", delivery.Body)
		}
		if !signature.IgnoreWhenTaskNotRegistered {
			delivery.Nack(multiple, requeue)
		}
		return nil
	}

	log.DEBUG.Printf("Received new message: %s", delivery.Body)

	err := taskProcessor.Process(signature)
	delivery.Ack(multiple)
	return err
}

// delay a task by delayDuration miliseconds, the way it works is a new queue
// is created without any consumers, the message is then published to this queue
// with appropriate ttl expiration headers, after the expiration, it is sent to
// the proper queue with consumers
func (b *Broker) delay(signature *tasks.Signature, delayMs int64) error {
	if delayMs <= 0 {
		return errors.New("Cannot delay task by 0ms")
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// It's necessary to redeclare the queue each time (to zero its TTL timer).
	queueName := fmt.Sprintf(
		"delay.%d.%s.%s",
		delayMs, // delay duration in mileseconds
		b.GetConfig().AMQP.Exchange,
		signature.RoutingKey, // routing key
	)
	declareQueueArgs := amqp.Table{
		// Exchange where to send messages after TTL expiration.
		"x-dead-letter-exchange": b.GetConfig().AMQP.Exchange,
		// Routing key which use when resending expired messages.
		"x-dead-letter-routing-key": signature.RoutingKey,
		// Time in milliseconds
		// after that message will expire and be sent to destination.
		"x-message-ttl": delayMs,
		// Time after that the queue will be deleted.
		"x-expires": delayMs * 2,
	}
	conn, channel, _, _, _, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().MultipleBrokerSeparator,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		b.GetConfig().AMQP.AutoDelete,   // queue delete when unused
		queueName,                       // queue binding key
		nil,                             // exchange declare args
		declareQueueArgs,                // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	defer b.Close(channel, conn)

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange
		queueName,                   // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	return nil
}

func (b *Broker) isDirectExchange() bool {
	return b.GetConfig().AMQP != nil && b.GetConfig().AMQP.ExchangeType == "direct"
}

// AdjustRoutingKey makes sure the routing key is correct.
// If the routing key is an empty string:
// a) set it to binding key for direct exchange type
// b) set it to default queue name
func (b *Broker) AdjustRoutingKey(s *tasks.Signature) {
	if s.RoutingKey != "" {
		return
	}

	if b.isDirectExchange() {
		// The routing algorithm behind a direct exchange is simple - a message goes
		// to the queues whose binding key exactly matches the routing key of the message.
		s.RoutingKey = b.GetConfig().AMQP.BindingKey
		return
	}

	s.RoutingKey = b.GetConfig().DefaultQueue
}
<commit_msg>Allow multiple requeues (#370)<commit_after>package amqp

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/RichardKnop/machinery/v1/brokers/errs"
	"github.com/RichardKnop/machinery/v1/brokers/iface"
	"github.com/RichardKnop/machinery/v1/common"
	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/pkg/errors"
	"github.com/streadway/amqp"
)

type AMQPConnection struct {
	queueName    string
	connection   *amqp.Connection
	channel      *amqp.Channel
	queue        amqp.Queue
	confirmation <-chan amqp.Confirmation
	errorchan    <-chan *amqp.Error
	cleanup      chan struct{}
}

// Broker represents an AMQP broker
type Broker struct {
	common.Broker
	common.AMQPConnector
	processingWG sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal

	connections      map[string]*AMQPConnection
	connectionsMutex sync.RWMutex
}

// New creates new Broker instance
func New(cnf *config.Config) iface.Broker {
	return &Broker{Broker: common.NewBroker(cnf), AMQPConnector: common.AMQPConnector{}, connections: make(map[string]*AMQPConnection)}
}

// StartConsuming enters a loop and waits for incoming messages
func (b *Broker) StartConsuming(consumerTag string, concurrency int, taskProcessor iface.TaskProcessor) (bool, error) {
	b.Broker.StartConsuming(consumerTag, concurrency, taskProcessor)

	queueName := taskProcessor.CustomQueue()
	if queueName == "" {
		queueName = b.GetConfig().DefaultQueue
	}

	conn, channel, queue, _, amqpCloseChan, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().MultipleBrokerSeparator,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		false,                           // queue delete when unused
		b.GetConfig().AMQP.BindingKey,   // queue binding key
		nil,                             // exchange declare args
		amqp.Table(b.GetConfig().AMQP.QueueDeclareArgs), // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		b.GetRetryFunc()(b.GetRetryStopChan())
		return b.GetRetry(), err
	}
	defer b.Close(channel, conn)

	if err = channel.Qos(
		b.GetConfig().AMQP.PrefetchCount,
		0,     // prefetch size
		false, // global
	); err != nil {
		return b.GetRetry(), fmt.Errorf("Channel qos error: %s", err)
	}

	deliveries, err := channel.Consume(
		queue.Name,  // queue
		consumerTag, // consumer tag
		false,       // auto-ack
		false,       // exclusive
		false,       // no-local
		false,       // no-wait
		nil,         // arguments
	)
	if err != nil {
		return b.GetRetry(), fmt.Errorf("Queue consume error: %s", err)
	}

	log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

	if err := b.consume(deliveries, concurrency, taskProcessor, amqpCloseChan); err != nil {
		return b.GetRetry(), err
	}

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	return b.GetRetry(), nil
}

// StopConsuming quits the loop
func (b *Broker) StopConsuming() {
	b.Broker.StopConsuming()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()
}

// GetOrOpenConnection will return a connection on a particular queue name. Open connections
// are saved to avoid having to reopen connection for multiple queues
func (b *Broker) GetOrOpenConnection(queueName string, queueBindingKey string, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs amqp.Table) (*AMQPConnection, error) {
	var err error

	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	conn, ok := b.connections[queueName]
	if !ok {
		conn = &AMQPConnection{
			queueName: queueName,
			cleanup:   make(chan struct{}),
		}
		conn.connection, conn.channel, conn.queue, conn.confirmation, conn.errorchan, err = b.Connect(
			b.GetConfig().Broker,
			b.GetConfig().MultipleBrokerSeparator,
			b.GetConfig().TLSConfig,
			b.GetConfig().AMQP.Exchange,     // exchange name
			b.GetConfig().AMQP.ExchangeType, // exchange type
			queueName,                       // queue name
			true,                            // queue durable
			false,                           // queue delete when unused
			queueBindingKey,                 // queue binding key
			exchangeDeclareArgs,             // exchange declare args
			queueDeclareArgs,                // queue declare args
			queueBindingArgs,                // queue binding args
		)
		if err != nil {
			return nil, errors.Wrapf(err, "Failed to connect to queue %s", queueName)
		}

		// Reconnect to the channel if it disconnects/errors out
		go func() {
			select {
			case err = <-conn.errorchan:
				log.INFO.Printf("Error occurred on queue: %s. Reconnecting", queueName)
				b.connectionsMutex.Lock()
				delete(b.connections, queueName)
				b.connectionsMutex.Unlock()
				_, err := b.GetOrOpenConnection(queueName, queueBindingKey, exchangeDeclareArgs, queueDeclareArgs, queueBindingArgs)
				if err != nil {
					log.ERROR.Printf("Failed to reopen queue: %s.", queueName)
				}
			case <-conn.cleanup:
				return
			}
			return
		}()
		b.connections[queueName] = conn
	}
	return conn, nil
}

func (b *Broker) CloseConnections() error {
	b.connectionsMutex.Lock()
	defer b.connectionsMutex.Unlock()

	for key, conn := range b.connections {
		if err := b.Close(conn.channel, conn.connection); err != nil {
			log.ERROR.Print("Failed to close channel")
			return nil
		}
		close(conn.cleanup)
		delete(b.connections, key)
	}
	return nil
}

// Publish places a new message on the default queue
func (b *Broker) Publish(ctx context.Context, signature *tasks.Signature) error {
	// Adjust routing key (this decides which queue the message will be published to)
	b.AdjustRoutingKey(signature)

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// delay the task
	if signature.ETA != nil {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			delayMs := int64(signature.ETA.Sub(now) / time.Millisecond)

			return b.delay(signature, delayMs)
		}
	}

	queue := b.GetConfig().DefaultQueue
	bindingKey := b.GetConfig().AMQP.BindingKey // queue binding key
	if b.isDirectExchange() {
		queue = signature.RoutingKey
		bindingKey = signature.RoutingKey
	}

	connection, err := b.GetOrOpenConnection(
		queue,
		bindingKey, // queue binding key
		nil,        // exchange declare args
		amqp.Table(b.GetConfig().AMQP.QueueDeclareArgs), // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return errors.Wrapf(err, "Failed to get a connection for queue %s", queue)
	}

	channel := connection.channel
	confirmsChan := connection.confirmation

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange name
		signature.RoutingKey,        // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         msg,
			Priority:     signature.Priority,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return errors.Wrap(err, "Failed to publish task")
	}

	confirmed := <-confirmsChan

	if confirmed.Ack {
		return nil
	}

	return fmt.Errorf("Failed delivery of delivery tag: %v", confirmed.DeliveryTag)
}

// consume takes delivered messages from the channel and manages a worker pool
// to process tasks concurrently
func (b *Broker) consume(deliveries <-chan amqp.Delivery, concurrency int, taskProcessor iface.TaskProcessor, amqpCloseChan <-chan *amqp.Error) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		for i := 0; i < concurrency; i++ {
			pool <- struct{}{}
		}
	}()

	// make channel with a capacity makes it become a buffered channel so that a worker which wants to
	// push an error to `errorsChan` doesn't need to be blocked while the for-loop is blocked waiting
	// a worker, that is, it avoids a possible deadlock
	errorsChan := make(chan error, 1)

	for {
		select {
		case amqpErr := <-amqpCloseChan:
			return amqpErr
		case err := <-errorsChan:
			return err
		case d := <-deliveries:
			if concurrency > 0 {
				// get worker from pool (blocks until one is available)
				<-pool
			}

			b.processingWG.Add(1)

			// Consume the task inside a gotourine so multiple tasks
			// can be processed concurrently
			go func() {
				if err := b.consumeOne(d, taskProcessor); err != nil {
					errorsChan <- err
				}

				b.processingWG.Done()

				if concurrency > 0 {
					// give worker back to pool
					pool <- struct{}{}
				}
			}()
		case <-b.GetStopChan():
			return nil
		}
	}
}

// consumeOne processes a single message using TaskProcessor
func (b *Broker) consumeOne(delivery amqp.Delivery, taskProcessor iface.TaskProcessor) error {
	if len(delivery.Body) == 0 {
		delivery.Nack(true, false)                     // multiple, requeue
		return errors.New("Received an empty message") // RabbitMQ down?
	}

	var multiple, requeue = false, false

	// Unmarshal message body into signature struct
	signature := new(tasks.Signature)
	decoder := json.NewDecoder(bytes.NewReader(delivery.Body))
	decoder.UseNumber()
	if err := decoder.Decode(signature); err != nil {
		delivery.Nack(multiple, requeue)
		return errs.NewErrCouldNotUnmarshaTaskSignature(delivery.Body, err)
	}

	// If the task is not registered, we nack it and requeue,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(signature.Name) {
		requeue = true
		log.INFO.Printf("Task not registered with this worker. Requeing message: %s", delivery.Body)

    if !signature.IgnoreWhenTaskNotRegistered {
			delivery.Nack(multiple, requeue)
		}

		return nil
	}

	log.DEBUG.Printf("Received new message: %s", delivery.Body)

	err := taskProcessor.Process(signature)
	delivery.Ack(multiple)
	return err
}

// delay a task by delayDuration miliseconds, the way it works is a new queue
// is created without any consumers, the message is then published to this queue
// with appropriate ttl expiration headers, after the expiration, it is sent to
// the proper queue with consumers
func (b *Broker) delay(signature *tasks.Signature, delayMs int64) error {
	if delayMs <= 0 {
		return errors.New("Cannot delay task by 0ms")
	}

	message, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// It's necessary to redeclare the queue each time (to zero its TTL timer).
	queueName := fmt.Sprintf(
		"delay.%d.%s.%s",
		delayMs, // delay duration in mileseconds
		b.GetConfig().AMQP.Exchange,
		signature.RoutingKey, // routing key
	)
	declareQueueArgs := amqp.Table{
		// Exchange where to send messages after TTL expiration.
		"x-dead-letter-exchange": b.GetConfig().AMQP.Exchange,
		// Routing key which use when resending expired messages.
		"x-dead-letter-routing-key": signature.RoutingKey,
		// Time in milliseconds
		// after that message will expire and be sent to destination.
		"x-message-ttl": delayMs,
		// Time after that the queue will be deleted.
		"x-expires": delayMs * 2,
	}
	conn, channel, _, _, _, err := b.Connect(
		b.GetConfig().Broker,
		b.GetConfig().MultipleBrokerSeparator,
		b.GetConfig().TLSConfig,
		b.GetConfig().AMQP.Exchange,     // exchange name
		b.GetConfig().AMQP.ExchangeType, // exchange type
		queueName,                       // queue name
		true,                            // queue durable
		b.GetConfig().AMQP.AutoDelete,   // queue delete when unused
		queueName,                       // queue binding key
		nil,                             // exchange declare args
		declareQueueArgs,                // queue declare args
		amqp.Table(b.GetConfig().AMQP.QueueBindingArgs), // queue binding args
	)
	if err != nil {
		return err
	}

	defer b.Close(channel, conn)

	if err := channel.Publish(
		b.GetConfig().AMQP.Exchange, // exchange
		queueName,                   // routing key
		false,                       // mandatory
		false,                       // immediate
		amqp.Publishing{
			Headers:      amqp.Table(signature.Headers),
			ContentType:  "application/json",
			Body:         message,
			DeliveryMode: amqp.Persistent,
		},
	); err != nil {
		return err
	}

	return nil
}

func (b *Broker) isDirectExchange() bool {
	return b.GetConfig().AMQP != nil && b.GetConfig().AMQP.ExchangeType == "direct"
}

// AdjustRoutingKey makes sure the routing key is correct.
// If the routing key is an empty string:
// a) set it to binding key for direct exchange type
// b) set it to default queue name
func (b *Broker) AdjustRoutingKey(s *tasks.Signature) {
	if s.RoutingKey != "" {
		return
	}

	if b.isDirectExchange() {
		// The routing algorithm behind a direct exchange is simple - a message goes
		// to the queues whose binding key exactly matches the routing key of the message.
		s.RoutingKey = b.GetConfig().AMQP.BindingKey
		return
	}

	s.RoutingKey = b.GetConfig().DefaultQueue
}
<|endoftext|>