<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package addressupdater

import (
	"fmt"
	"time"

	"launchpad.net/loggo"

	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/watcher"
)

var logger = loggo.GetLogger("juju.worker.addressupdater")

// ShortPoll and LongPoll hold the polling intervales for the
// address updater. When a machine has no address,
// it will be polled at ShortPoll intervals until it does;
// after that LongPoll will be used to check that
// the instance address has not changed.
var (
	ShortPoll = 1 * time.Second
	LongPoll  = 1 * time.Minute
)

type machine interface {
	Id() string
	Addresses() []instance.Address
	InstanceId() (instance.Id, error)
	SetAddresses([]instance.Address) error
	String() string
	Refresh() error
	Life() state.Life
}

type machineContext interface {
	killAll(err error)
	addresses(id instance.Id) ([]instance.Address, error)
	dying() <-chan struct{}
}

type machineAddress struct {
	machine   machine
	addresses []instance.Address
}

var _ machine = (*state.Machine)(nil)

type machinesWatcher interface {
	Changes() <-chan []string
	Err() error
	Stop() error
}

type updaterContext interface {
	newMachineContext() machineContext
	getMachine(id string) (machine, error)
	dying() <-chan struct{}
}

type updater struct {
	context     updaterContext
	machines    map[string]chan struct{}
	machineDead chan machine
}

// watchMachinesLoop watches for changes provided by the given
// machinesWatcher and starts machine goroutines to deal
// with them, using the provided newMachineContext
// function to create the appropriate context for each new machine id.
func watchMachinesLoop(context updaterContext, w machinesWatcher) (err error) {
	p := &updater{
		context:     context,
		machines:    make(map[string]chan struct{}),
		machineDead: make(chan machine),
	}
	defer func() {
		if stopErr := w.Stop(); stopErr != nil {
			if err == nil {
				err = fmt.Errorf("error stopping watcher: %v", stopErr)
			} else {
				logger.Warningf("ignoring error when stopping watcher: %v", stopErr)
			}
		}
		for len(p.machines) > 0 {
			delete(p.machines, (<-p.machineDead).Id())
		}
	}()
	for {
		select {
		case ids, ok := <-w.Changes():
			if !ok {
				return watcher.MustErr(w)
			}
			if err := p.startMachines(ids); err != nil {
				return err
			}
		case m := <-p.machineDead:
			delete(p.machines, m.Id())
		case <-p.context.dying():
			return nil
		}
	}
}

func (p *updater) startMachines(ids []string) error {
	for _, id := range ids {
		if c := p.machines[id]; c == nil {
			// We don't know about the machine - start
			// a goroutine to deal with it.
			m, err := p.context.getMachine(id)
			if errors.IsNotFoundError(err) {
				logger.Warningf("watcher gave notification of non-existent machine %q", id)
				continue
			}
			if err != nil {
				return err
			}
			c = make(chan struct{})
			p.machines[id] = c
			go runMachine(p.context.newMachineContext(), m, c, p.machineDead)
		} else {
			c <- struct{}{}
		}
	}
	return nil
}

// runMachine processes the address publishing for a given machine.
// We assume that the machine is alive when this is first called.
func runMachine(context machineContext, m machine, changed <-chan struct{}, died chan<- machine) {
	defer func() {
		// We can't just send on the died channel because the
		// central loop might be trying to write to us on the
		// changed channel.
		for {
			select {
			case died <- m:
				return
			case <-changed:
			}
		}
	}()
	if err := machineLoop(context, m, changed); err != nil {
		context.killAll(err)
	}
}

func machineLoop(context machineContext, m machine, changed <-chan struct{}) error {
	// Use a short poll interval when initially waiting for
	// a machine's address, and a long one when it already
	// has an address.
	pollInterval := ShortPoll
	checkAddress := true
	for {
		if checkAddress {
			if err := checkMachineAddresses(context, m); err != nil {
				if errors.IsNotImplementedError(err) {
					pollInterval = 365 * 24 * time.Hour
				}
				return err
			}
			if len(m.Addresses()) > 0 {
				pollInterval = LongPoll
			}
			checkAddress = false
		}
		select {
		case <-time.After(pollInterval):
			checkAddress = true
		case <-context.dying():
			return nil
		case <-changed:
			if err := m.Refresh(); err != nil {
				return err
			}
			if m.Life() == state.Dead {
				return nil
			}
		}
	}
}

// checkMachineAddresses checks the current provider addresses
// for the given machine's instance, and sets them
// on the machine if they've changed.
func checkMachineAddresses(context machineContext, m machine) error {
	instId, err := m.InstanceId()
	if err != nil && !state.IsNotProvisionedError(err) {
		return fmt.Errorf("cannot get machine's instance id: %v", err)
	}
	var newAddrs []instance.Address
	if err == nil {
		newAddrs, err = context.addresses(instId)
		if err != nil {
			if errors.IsNotImplementedError(err) {
				return err
			}
			logger.Warningf("cannot get addresses for instance %q: %v", instId, err)
			return nil
		}
	}
	if addressesEqual(m.Addresses(), newAddrs) {
		return nil
	}
	if err := m.SetAddresses(newAddrs); err != nil {
		return fmt.Errorf("cannot set addresses on %q: %v", m, err)
	}
	return nil
}

func addressesEqual(a0, a1 []instance.Address) bool {
	if len(a0) != len(a1) {
		return false
	}
	for i := range a0 {
		if a0[i] != a1[i] {
			return false
		}
	}
	return true
}
<commit_msg>worker/addressupdater: add comment<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package addressupdater

import (
	"fmt"
	"time"

	"launchpad.net/loggo"

	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/watcher"
)

var logger = loggo.GetLogger("juju.worker.addressupdater")

// ShortPoll and LongPoll hold the polling intervales for the
// address updater. When a machine has no address,
// it will be polled at ShortPoll intervals until it does;
// after that LongPoll will be used to check that
// the instance address has not changed.
var (
	ShortPoll = 1 * time.Second
	LongPoll  = 1 * time.Minute
)

type machine interface {
	Id() string
	Addresses() []instance.Address
	InstanceId() (instance.Id, error)
	SetAddresses([]instance.Address) error
	String() string
	Refresh() error
	Life() state.Life
}

type machineContext interface {
	killAll(err error)
	addresses(id instance.Id) ([]instance.Address, error)
	dying() <-chan struct{}
}

type machineAddress struct {
	machine   machine
	addresses []instance.Address
}

var _ machine = (*state.Machine)(nil)

type machinesWatcher interface {
	Changes() <-chan []string
	Err() error
	Stop() error
}

type updaterContext interface {
	newMachineContext() machineContext
	getMachine(id string) (machine, error)
	dying() <-chan struct{}
}

type updater struct {
	context     updaterContext
	machines    map[string]chan struct{}
	machineDead chan machine
}

// watchMachinesLoop watches for changes provided by the given
// machinesWatcher and starts machine goroutines to deal
// with them, using the provided newMachineContext
// function to create the appropriate context for each new machine id.
func watchMachinesLoop(context updaterContext, w machinesWatcher) (err error) {
	p := &updater{
		context:     context,
		machines:    make(map[string]chan struct{}),
		machineDead: make(chan machine),
	}
	defer func() {
		if stopErr := w.Stop(); stopErr != nil {
			if err == nil {
				err = fmt.Errorf("error stopping watcher: %v", stopErr)
			} else {
				logger.Warningf("ignoring error when stopping watcher: %v", stopErr)
			}
		}
		for len(p.machines) > 0 {
			delete(p.machines, (<-p.machineDead).Id())
		}
	}()
	for {
		select {
		case ids, ok := <-w.Changes():
			if !ok {
				return watcher.MustErr(w)
			}
			if err := p.startMachines(ids); err != nil {
				return err
			}
		case m := <-p.machineDead:
			delete(p.machines, m.Id())
		case <-p.context.dying():
			return nil
		}
	}
}

func (p *updater) startMachines(ids []string) error {
	for _, id := range ids {
		if c := p.machines[id]; c == nil {
			// We don't know about the machine - start
			// a goroutine to deal with it.
			m, err := p.context.getMachine(id)
			if errors.IsNotFoundError(err) {
				logger.Warningf("watcher gave notification of non-existent machine %q", id)
				continue
			}
			if err != nil {
				return err
			}
			c = make(chan struct{})
			p.machines[id] = c
			go runMachine(p.context.newMachineContext(), m, c, p.machineDead)
		} else {
			c <- struct{}{}
		}
	}
	return nil
}

// runMachine processes the address publishing for a given machine.
// We assume that the machine is alive when this is first called.
func runMachine(context machineContext, m machine, changed <-chan struct{}, died chan<- machine) {
	defer func() {
		// We can't just send on the died channel because the
		// central loop might be trying to write to us on the
		// changed channel.
		for {
			select {
			case died <- m:
				return
			case <-changed:
			}
		}
	}()
	if err := machineLoop(context, m, changed); err != nil {
		context.killAll(err)
	}
}

func machineLoop(context machineContext, m machine, changed <-chan struct{}) error {
	// Use a short poll interval when initially waiting for
	// a machine's address, and a long one when it already
	// has an address.
	pollInterval := ShortPoll
	checkAddress := true
	for {
		if checkAddress {
			if err := checkMachineAddresses(context, m); err != nil {
				// If the provider doesn't implement addresses now,
				// it never will until we're upgraded, so don't bother
				// asking any more. We could use less resources
				// by taking down the entire address updater worker,
				// but this is easier for now (and hopefully the local
				// provider will implement Addresses in the not-too-distant
				// future), so we won't need to worry about this case at all.
				if errors.IsNotImplementedError(err) {
					pollInterval = 365 * 24 * time.Hour
				}
				return err
			}
			if len(m.Addresses()) > 0 {
				pollInterval = LongPoll
			}
			checkAddress = false
		}
		select {
		case <-time.After(pollInterval):
			checkAddress = true
		case <-context.dying():
			return nil
		case <-changed:
			if err := m.Refresh(); err != nil {
				return err
			}
			if m.Life() == state.Dead {
				return nil
			}
		}
	}
}

// checkMachineAddresses checks the current provider addresses
// for the given machine's instance, and sets them
// on the machine if they've changed.
func checkMachineAddresses(context machineContext, m machine) error {
	instId, err := m.InstanceId()
	if err != nil && !state.IsNotProvisionedError(err) {
		return fmt.Errorf("cannot get machine's instance id: %v", err)
	}
	var newAddrs []instance.Address
	if err == nil {
		newAddrs, err = context.addresses(instId)
		if err != nil {
			if errors.IsNotImplementedError(err) {
				return err
			}
			logger.Warningf("cannot get addresses for instance %q: %v", instId, err)
			return nil
		}
	}
	if addressesEqual(m.Addresses(), newAddrs) {
		return nil
	}
	if err := m.SetAddresses(newAddrs); err != nil {
		return fmt.Errorf("cannot set addresses on %q: %v", m, err)
	}
	return nil
}

func addressesEqual(a0, a1 []instance.Address) bool {
	if len(a0) != len(a1) {
		return false
	}
	for i := range a0 {
		if a0[i] != a1[i] {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package shaderManager

import (
	"fmt"
	"io/ioutil"
	"strings"
	"sync"

	"github.com/go-gl/gl/v4.1-core/gl"
)

//Shader holds information about a shader program
type Shader struct {
	VertSrcFile string
	FragSrcFile string
	VertSrc     string
	FragSrc     string
	Name        string
}

//ShaderManager stores shader programs
type shaderManager struct {
	programs      map[string]uint32
	programLock   sync.RWMutex
	DefaultShader string
}

//ShaderManager interface is used to interact with the shaderManager
type ShaderManager interface {
	LoadProgram(Shader, bool)
	GetShader(key string) (uint32, bool)
	GetDefaultShader() string
}

//NewShaderManager creates a new ShaderManager
func NewShaderManager() ShaderManager {
	sm := shaderManager{programs: make(map[string]uint32)}
	return &sm
}

//LoadProgram creates a shader program from a vertex and fragment shader source files.
func (sm *shaderManager) LoadProgram(shader Shader, shouldBeDefault bool) {

	if len(shader.VertSrc) < 1 {
		simpleVert, err := loadShader(shader.VertSrcFile)
		shader.VertSrc = simpleVert
		if err != nil {
			fmt.Println(err)
			return
		}
	}

	if len(shader.FragSrc) < 1 {
		simpleFrag, err := loadShader(shader.FragSrcFile)
		shader.FragSrc = simpleFrag
		if err != nil {
			fmt.Println(err)
			return
		}
	}

	program, err := newProgram(shader.VertSrc, shader.FragSrc)
	if err != nil {
		fmt.Println(err)
		return
	}

	sm.programLock.Lock()
	defer sm.programLock.Unlock()
	sm.programs[shader.Name] = program

	if len(sm.programs) == 1 || shouldBeDefault {
		sm.DefaultShader = shader.Name
	}
}

//GetShader returns a program id if the shader program was loaded, if it was not a 0
//false will be returned
func (sm *shaderManager) GetShader(key string) (uint32, bool) {

	sm.programLock.RLock()
	defer sm.programLock.RUnlock()
	program, status := sm.programs[key]

	return program, status
}

//GetDefaultShader returns the name of the default shader
func (sm *shaderManager) GetDefaultShader() string {
	return sm.DefaultShader
}

func newProgram(vertexShaderSource, fragmentShaderSource string) (uint32, error) {
	vertexShader, err := compileShader(vertexShaderSource, gl.VERTEX_SHADER)
	if err != nil {
		return 0, err
	}

	fragmentShader, err := compileShader(fragmentShaderSource, gl.FRAGMENT_SHADER)
	if err != nil {
		return 0, err
	}

	program := gl.CreateProgram()

	gl.AttachShader(program, vertexShader)
	gl.AttachShader(program, fragmentShader)
	gl.LinkProgram(program)

	var status int32
	gl.GetProgramiv(program, gl.LINK_STATUS, &status)
	if status == gl.FALSE {
		var logLength int32
		gl.GetProgramiv(program, gl.INFO_LOG_LENGTH, &logLength)

		log := strings.Repeat("\x00", int(logLength+1))
		gl.GetProgramInfoLog(program, logLength, nil, gl.Str(log))

		return 0, fmt.Errorf("failed to link program: %v", log)
	}

	gl.DeleteShader(vertexShader)
	gl.DeleteShader(fragmentShader)

	return program, nil
}

func compileShader(source string, shaderType uint32) (uint32, error) {
	shader := gl.CreateShader(shaderType)

	csource := gl.Str(source)
	gl.ShaderSource(shader, 1, &csource, nil)
	gl.CompileShader(shader)

	var status int32
	gl.GetShaderiv(shader, gl.COMPILE_STATUS, &status)
	if status == gl.FALSE {
		var logLength int32
		gl.GetShaderiv(shader, gl.INFO_LOG_LENGTH, &logLength)

		log := strings.Repeat("\x00", int(logLength+1))
		gl.GetShaderInfoLog(shader, logLength, nil, gl.Str(log))

		return 0, fmt.Errorf("failed to compile %v: %v", source, log)
	}

	return shader, nil
}

func loadShader(fileName string) (string, error) {
	data, err := ioutil.ReadFile(fileName)
	if err != nil {
		return "", err
	}

	return string(data) + "\x00", nil
}
<commit_msg>Added check to make sure shader src has correct suffix<commit_after>package shaderManager

import (
	"fmt"
	"io/ioutil"
	"strings"
	"sync"

	"github.com/go-gl/gl/v4.1-core/gl"
)

//Shader holds information about a shader program
//VertSrcFile file that contains the vertex shader
//FragSrcFile file that contains the fragment shader
//VertSrc: vertex shader source
//FragSrc: fragment shader source
type Shader struct {
	VertSrcFile string
	FragSrcFile string
	VertSrc     string
	FragSrc     string
	Name        string
}

//shaderManager stores shader programs
type shaderManager struct {
	programs      map[string]uint32
	programLock   sync.RWMutex
	DefaultShader string
}

//ShaderManager interface is used to interact with the shaderManager
type ShaderManager interface {
	LoadProgram(Shader, bool)
	GetShader(key string) (uint32, bool)
	GetDefaultShader() string
}

//NewShaderManager creates a new ShaderManager
func NewShaderManager() ShaderManager {
	sm := shaderManager{programs: make(map[string]uint32)}
	return &sm
}

//LoadProgram creates a shader program from a vertex and fragment shader source files.
func (sm *shaderManager) LoadProgram(shader Shader, shouldBeDefault bool) {

	if len(shader.VertSrc) < 1 {
		simpleVert, err := loadShader(shader.VertSrcFile)
		shader.VertSrc = simpleVert
		if err != nil {
			fmt.Println(err)
			return
		}
	} else if !strings.HasSuffix(shader.VertSrc, "\x00") {
		shader.VertSrc = shader.VertSrc + "\x00"
	}

	if len(shader.FragSrc) < 1 {
		simpleFrag, err := loadShader(shader.FragSrcFile)
		shader.FragSrc = simpleFrag
		if err != nil {
			fmt.Println(err)
			return
		}
	} else if !strings.HasSuffix(shader.FragSrc, "\x00") {
		shader.FragSrc = shader.FragSrc + "\x00"
	}

	program, err := newProgram(shader.VertSrc, shader.FragSrc)
	if err != nil {
		fmt.Println(err)
		return
	}

	sm.programLock.Lock()
	defer sm.programLock.Unlock()
	sm.programs[shader.Name] = program

	if len(sm.programs) == 1 || shouldBeDefault {
		sm.DefaultShader = shader.Name
	}
}

//GetShader returns a program id if the shader program was loaded, if it was not a 0
//false will be returned
func (sm *shaderManager) GetShader(key string) (uint32, bool) {

	sm.programLock.RLock()
	defer sm.programLock.RUnlock()
	program, status := sm.programs[key]

	return program, status
}

//GetDefaultShader returns the name of the default shader
func (sm *shaderManager) GetDefaultShader() string {
	return sm.DefaultShader
}

func newProgram(vertexShaderSource, fragmentShaderSource string) (uint32, error) {
	vertexShader, err := compileShader(vertexShaderSource, gl.VERTEX_SHADER)
	if err != nil {
		return 0, err
	}

	fragmentShader, err := compileShader(fragmentShaderSource, gl.FRAGMENT_SHADER)
	if err != nil {
		return 0, err
	}

	program := gl.CreateProgram()

	gl.AttachShader(program, vertexShader)
	gl.AttachShader(program, fragmentShader)
	gl.LinkProgram(program)

	var status int32
	gl.GetProgramiv(program, gl.LINK_STATUS, &status)
	if status == gl.FALSE {
		var logLength int32
		gl.GetProgramiv(program, gl.INFO_LOG_LENGTH, &logLength)

		log := strings.Repeat("\x00", int(logLength+1))
		gl.GetProgramInfoLog(program, logLength, nil, gl.Str(log))

		return 0, fmt.Errorf("failed to link program: %v", log)
	}

	gl.DeleteShader(vertexShader)
	gl.DeleteShader(fragmentShader)

	return program, nil
}

func compileShader(source string, shaderType uint32) (uint32, error) {
	shader := gl.CreateShader(shaderType)

	csource := gl.Str(source)
	gl.ShaderSource(shader, 1, &csource, nil)
	gl.CompileShader(shader)

	var status int32
	gl.GetShaderiv(shader, gl.COMPILE_STATUS, &status)
	if status == gl.FALSE {
		var logLength int32
		gl.GetShaderiv(shader, gl.INFO_LOG_LENGTH, &logLength)

		log := strings.Repeat("\x00", int(logLength+1))
		gl.GetShaderInfoLog(shader, logLength, nil, gl.Str(log))

		return 0, fmt.Errorf("failed to compile %v: %v", source, log)
	}

	return shader, nil
}

func loadShader(fileName string) (string, error) {
	data, err := ioutil.ReadFile(fileName)
	if err != nil {
		return "", err
	}

	return string(data) + "\x00", nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/rcrowley/go-metrics"
	"github.com/skia-dev/glog"
	"go.skia.org/infra/alertserver/go/alerting"
	"go.skia.org/infra/go/autoroll"
	"go.skia.org/infra/go/buildbot"
	"go.skia.org/infra/go/influxdb"
	"go.skia.org/infra/go/util"
)

/*
	This file contains goroutines which trigger more complex alerts than
	can be expressed using the rule format in alerts.cfg.
*/

const (
	ANDROID_DISCONNECT = `The Android device for %s appears to be disconnected.

Build: https://uberchromegw.corp.google.com/i/%s/builders/%s/builds/%d
Dashboard: https://status.skia.org/buildbots?botGrouping=buildslave&filterBy=buildslave&include=%%5E%s%%24&tab=builds
Host info: https://status.skia.org/hosts?filter=%s`
	AUTOROLL_ALERT_NAME = "AutoRoll Failed"
	BUILDSLAVE_OFFLINE  = `Buildslave %s is not connected to https://uberchromegw.corp.google.com/i/%s/buildslaves/%s

Dashboard: https://status.skia.org/buildbots?botGrouping=buildslave&filterBy=buildslave&include=%%5E%s%%24&tab=builds
Host info: https://status.skia.org/hosts?filter=%s`
	HUNG_BUILDSLAVE = `Possibly hung buildslave (%s)

A step has been running for over %s:
https://uberchromegw.corp.google.com/i/%s/builders/%s/builds/%d
Dashboard: https://status.skia.org/buildbots?botGrouping=buildslave&filterBy=buildslave&include=%%5E%s%%24&tab=builds
Host info: https://status.skia.org/hosts?filter=%s`
	UPDATE_SCRIPTS = `update_scripts failed on %s

Build: https://uberchromegw.corp.google.com/i/%s/builders/%s/builds/%d
Dashboard: https://status.skia.org/buildbots?botGrouping=builder&filterBy=builder&include=%%5E%s%%24&tab=builds
Host info: https://status.skia.org/hosts?filter=%s`
)

var BUILDSLAVE_OFFLINE_BLACKLIST = []string{
	"build3-a3",
	"build4-a3",
	"vm255-m3",
}

type BuildSlice []*buildbot.Build

func (s BuildSlice) Len() int {
	return len(s)
}

func (s BuildSlice) Less(i, j int) bool {
	return s[i].Finished < s[j].Finished
}

func (s BuildSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func StartAlertRoutines(am *alerting.AlertManager, tickInterval time.Duration, c *influxdb.Client) {
	emailAction, err := alerting.ParseAction("Email(infra-alerts@skia.org)")
	if err != nil {
		glog.Fatal(err)
	}
	actions := []alerting.Action{emailAction}

	// Disconnected buildslaves.
	go func() {
		seriesTmpl := "buildbot.buildslaves.%s.connected"
		re := regexp.MustCompile("[^A-Za-z0-9]+")
		for _ = range time.Tick(tickInterval) {
			glog.Info("Loading buildslave data.")
			slaves, err := buildbot.GetBuildSlaves()
			if err != nil {
				glog.Error(err)
				continue
			}
			for masterName, m := range slaves {
				for _, s := range m {
					if util.In(s.Name, BUILDSLAVE_OFFLINE_BLACKLIST) {
						continue
					}
					v := int64(0)
					if s.Connected {
						v = int64(1)
					}
					metric := fmt.Sprintf(seriesTmpl, re.ReplaceAllString(s.Name, "_"))
					metrics.GetOrRegisterGauge(metric, metrics.DefaultRegistry).Update(v)
					if !s.Connected {
						// This buildslave is offline. Figure out which one it is.
						if err := am.AddAlert(&alerting.Alert{
							Name:        fmt.Sprintf("Buildslave %s offline", s.Name),
							Category:    alerting.INFRA_ALERT,
							Message:     fmt.Sprintf(BUILDSLAVE_OFFLINE, s.Name, masterName, s.Name, s.Name, s.Name),
							Nag:         int64(time.Hour),
							AutoDismiss: int64(2 * tickInterval),
							Actions:     actions,
						}); err != nil {
							glog.Error(err)
						}
					}
				}
			}
		}
	}()

	// AutoRoll failure.
	go func() {
		lastSearch := time.Now()
		for now := range time.Tick(time.Minute) {
			glog.Infof("Searching for DEPS rolls.")
			results, err := autoroll.GetRecentRolls(lastSearch)
			if err != nil {
				glog.Errorf("Failed to search for DEPS rolls: %v", err)
				continue
			}
			lastSearch = now
			activeAlert := am.ActiveAlert(AUTOROLL_ALERT_NAME)
			for _, issue := range results {
				if issue.Closed {
					if issue.Committed {
						if activeAlert != 0 {
							msg := fmt.Sprintf("Subsequent roll succeeded: %s/%d", autoroll.RIETVELD_URL, issue.Issue)
							if err := am.Dismiss(activeAlert, alerting.USER_ALERTSERVER, msg); err != nil {
								glog.Error(err)
							}
						}
					} else {
						if err := am.AddAlert(&alerting.Alert{
							Name:    AUTOROLL_ALERT_NAME,
							Message: fmt.Sprintf("DEPS roll failed: %s/%d", autoroll.RIETVELD_URL, issue.Issue),
							Nag:     int64(3 * time.Hour),
							Actions: actions,
						}); err != nil {
							glog.Error(err)
						}
					}
				}
			}
		}
	}()

	// Android device disconnects, hung buildslaves.
	go func() {
		// These builders are frequently slow. Ignore them when looking for hung buildslaves.
		hungSlavesIgnore := []string{
			"Housekeeper-Nightly-RecreateSKPs_Canary",
			"Housekeeper-Weekly-RecreateSKPs",
			"Linux Builder",
			"Mac Builder",
			"Test-Ubuntu-GCC-GCE-CPU-AVX2-x86_64-Release-Valgrind",
			"Test-Ubuntu-GCC-ShuttleA-GPU-GTX550Ti-x86_64-Release-Valgrind",
			"Win Builder",
		}
		hangTimePeriod := 2 * time.Hour
		for _ = range time.Tick(tickInterval) {
			glog.Infof("Searching for hung buildslaves and disconnected Android devices.")
			builds, err := buildbot.GetUnfinishedBuilds()
			if err != nil {
				glog.Error(err)
				continue
			}
			for _, b := range builds {
				// Disconnected Android device?
				disconnectedAndroid := false
				if strings.Contains(b.Builder, "Android") && !strings.Contains(b.Builder, "Build") {
					for _, s := range b.Steps {
						if strings.Contains(s.Name, "wait for device") {
							// If "wait for device" has been running for 10 minutes, the device is probably offline.
							if s.Finished == 0 && time.Since(time.Unix(int64(s.Started), 0)) > 10*time.Minute {
								if err := am.AddAlert(&alerting.Alert{
									Name:     fmt.Sprintf("Android device disconnected (%s)", b.BuildSlave),
									Category: alerting.INFRA_ALERT,
									Message:  fmt.Sprintf(ANDROID_DISCONNECT, b.BuildSlave, b.Master, b.Builder, b.Number, b.BuildSlave, b.BuildSlave),
									Nag:      int64(3 * time.Hour),
									Actions:  actions,
								}); err != nil {
									glog.Error(err)
								}
								disconnectedAndroid = true
							}
						}
					}
				}
				if !disconnectedAndroid && !util.ContainsAny(b.Builder, hungSlavesIgnore) {
					// Hung buildslave?
					for _, s := range b.Steps {
						if s.Name == "steps" {
							continue
						}
						// If the step has been running for over an hour, it's probably hung.
						if s.Finished == 0 && time.Since(time.Unix(int64(s.Started), 0)) > hangTimePeriod {
							if err := am.AddAlert(&alerting.Alert{
								Name:        fmt.Sprintf("Possibly hung buildslave (%s)", b.BuildSlave),
								Category:    alerting.INFRA_ALERT,
								Message:     fmt.Sprintf(HUNG_BUILDSLAVE, b.BuildSlave, hangTimePeriod.String(), b.Master, b.Builder, b.Number, b.BuildSlave, b.BuildSlave),
								Nag:         int64(time.Hour),
								Actions:     actions,
								AutoDismiss: int64(10 * tickInterval),
							}); err != nil {
								glog.Error(err)
							}
						}
					}
				}
			}
		}
	}()

	// Failed update_scripts.
	go func() {
		lastSearch := time.Now()
		for _ = range time.Tick(tickInterval) {
			glog.Infof("Searching for builds which failed update_scripts.")
			currentSearch := time.Now()
			builds, err := buildbot.GetBuildsFromDateRange(lastSearch, currentSearch)
			lastSearch = currentSearch
			if err != nil {
				glog.Error(err)
				continue
			}
			for _, b := range builds {
				for _, s := range b.Steps {
					if s.Name == "update_scripts" {
						if s.Results != 0 {
							if err := am.AddAlert(&alerting.Alert{
								Name:     "update_scripts failed",
								Category: alerting.INFRA_ALERT,
								Message:  fmt.Sprintf(UPDATE_SCRIPTS, b.Builder, b.Master, b.Builder, b.Number, b.Builder, b.BuildSlave),
								Actions:  actions,
							}); err != nil {
								glog.Error(err)
							}
						}
						break
					}
				}
			}
		}
	}()
}
<commit_msg>Allow 3 hours for build steps<commit_after>package main

import (
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/rcrowley/go-metrics"
	"github.com/skia-dev/glog"
	"go.skia.org/infra/alertserver/go/alerting"
	"go.skia.org/infra/go/autoroll"
	"go.skia.org/infra/go/buildbot"
	"go.skia.org/infra/go/influxdb"
	"go.skia.org/infra/go/util"
)

/*
	This file contains goroutines which trigger more complex alerts than
	can be expressed using the rule format in alerts.cfg.
*/

const (
	ANDROID_DISCONNECT = `The Android device for %s appears to be disconnected.

Build: https://uberchromegw.corp.google.com/i/%s/builders/%s/builds/%d
Dashboard: https://status.skia.org/buildbots?botGrouping=buildslave&filterBy=buildslave&include=%%5E%s%%24&tab=builds
Host info: https://status.skia.org/hosts?filter=%s`
	AUTOROLL_ALERT_NAME = "AutoRoll Failed"
	BUILDSLAVE_OFFLINE  = `Buildslave %s is not connected to https://uberchromegw.corp.google.com/i/%s/buildslaves/%s

Dashboard: https://status.skia.org/buildbots?botGrouping=buildslave&filterBy=buildslave&include=%%5E%s%%24&tab=builds
Host info: https://status.skia.org/hosts?filter=%s`
	HUNG_BUILDSLAVE = `Possibly hung buildslave (%s)

A step has been running for over %s:
https://uberchromegw.corp.google.com/i/%s/builders/%s/builds/%d
Dashboard: https://status.skia.org/buildbots?botGrouping=buildslave&filterBy=buildslave&include=%%5E%s%%24&tab=builds
Host info: https://status.skia.org/hosts?filter=%s`
	UPDATE_SCRIPTS = `update_scripts failed on %s

Build: https://uberchromegw.corp.google.com/i/%s/builders/%s/builds/%d
Dashboard: https://status.skia.org/buildbots?botGrouping=builder&filterBy=builder&include=%%5E%s%%24&tab=builds
Host info: https://status.skia.org/hosts?filter=%s`
)

var BUILDSLAVE_OFFLINE_BLACKLIST = []string{
	"build3-a3",
	"build4-a3",
	"vm255-m3",
}

type BuildSlice []*buildbot.Build

func (s BuildSlice) Len() int {
	return len(s)
}

func (s BuildSlice) Less(i, j int) bool {
	return s[i].Finished < s[j].Finished
}

func (s BuildSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func StartAlertRoutines(am *alerting.AlertManager, tickInterval time.Duration, c *influxdb.Client) {
	emailAction, err := alerting.ParseAction("Email(infra-alerts@skia.org)")
	if err != nil {
		glog.Fatal(err)
	}
	actions := []alerting.Action{emailAction}

	// Disconnected buildslaves.
	go func() {
		seriesTmpl := "buildbot.buildslaves.%s.connected"
		re := regexp.MustCompile("[^A-Za-z0-9]+")
		for _ = range time.Tick(tickInterval) {
			glog.Info("Loading buildslave data.")
			slaves, err := buildbot.GetBuildSlaves()
			if err != nil {
				glog.Error(err)
				continue
			}
			for masterName, m := range slaves {
				for _, s := range m {
					if util.In(s.Name, BUILDSLAVE_OFFLINE_BLACKLIST) {
						continue
					}
					v := int64(0)
					if s.Connected {
						v = int64(1)
					}
					metric := fmt.Sprintf(seriesTmpl, re.ReplaceAllString(s.Name, "_"))
					metrics.GetOrRegisterGauge(metric, metrics.DefaultRegistry).Update(v)
					if !s.Connected {
						// This buildslave is offline. Figure out which one it is.
						if err := am.AddAlert(&alerting.Alert{
							Name:        fmt.Sprintf("Buildslave %s offline", s.Name),
							Category:    alerting.INFRA_ALERT,
							Message:     fmt.Sprintf(BUILDSLAVE_OFFLINE, s.Name, masterName, s.Name, s.Name, s.Name),
							Nag:         int64(time.Hour),
							AutoDismiss: int64(2 * tickInterval),
							Actions:     actions,
						}); err != nil {
							glog.Error(err)
						}
					}
				}
			}
		}
	}()

	// AutoRoll failure.
	go func() {
		lastSearch := time.Now()
		for now := range time.Tick(time.Minute) {
			glog.Infof("Searching for DEPS rolls.")
			results, err := autoroll.GetRecentRolls(lastSearch)
			if err != nil {
				glog.Errorf("Failed to search for DEPS rolls: %v", err)
				continue
			}
			lastSearch = now
			activeAlert := am.ActiveAlert(AUTOROLL_ALERT_NAME)
			for _, issue := range results {
				if issue.Closed {
					if issue.Committed {
						if activeAlert != 0 {
							msg := fmt.Sprintf("Subsequent roll succeeded: %s/%d", autoroll.RIETVELD_URL, issue.Issue)
							if err := am.Dismiss(activeAlert, alerting.USER_ALERTSERVER, msg); err != nil {
								glog.Error(err)
							}
						}
					} else {
						if err := am.AddAlert(&alerting.Alert{
							Name:    AUTOROLL_ALERT_NAME,
							Message: fmt.Sprintf("DEPS roll failed: %s/%d", autoroll.RIETVELD_URL, issue.Issue),
							Nag:     int64(3 * time.Hour),
							Actions: actions,
						}); err != nil {
							glog.Error(err)
						}
					}
				}
			}
		}
	}()

	// Android device disconnects, hung buildslaves.
	go func() {
		// These builders are frequently slow. Ignore them when looking for hung buildslaves.
		hungSlavesIgnore := []string{
			"Housekeeper-Nightly-RecreateSKPs_Canary",
			"Housekeeper-Weekly-RecreateSKPs",
			"Linux Builder",
			"Mac Builder",
			"Test-Ubuntu-GCC-GCE-CPU-AVX2-x86_64-Release-Valgrind",
			"Test-Ubuntu-GCC-ShuttleA-GPU-GTX550Ti-x86_64-Release-Valgrind",
			"Win Builder",
		}
		hangTimePeriod := 3 * time.Hour
		for _ = range time.Tick(tickInterval) {
			glog.Infof("Searching for hung buildslaves and disconnected Android devices.")
			builds, err := buildbot.GetUnfinishedBuilds()
			if err != nil {
				glog.Error(err)
				continue
			}
			for _, b := range builds {
				// Disconnected Android device?
				disconnectedAndroid := false
				if strings.Contains(b.Builder, "Android") && !strings.Contains(b.Builder, "Build") {
					for _, s := range b.Steps {
						if strings.Contains(s.Name, "wait for device") {
							// If "wait for device" has been running for 10 minutes, the device is probably offline.
							if s.Finished == 0 && time.Since(time.Unix(int64(s.Started), 0)) > 10*time.Minute {
								if err := am.AddAlert(&alerting.Alert{
									Name:     fmt.Sprintf("Android device disconnected (%s)", b.BuildSlave),
									Category: alerting.INFRA_ALERT,
									Message:  fmt.Sprintf(ANDROID_DISCONNECT, b.BuildSlave, b.Master, b.Builder, b.Number, b.BuildSlave, b.BuildSlave),
									Nag:      int64(3 * time.Hour),
									Actions:  actions,
								}); err != nil {
									glog.Error(err)
								}
								disconnectedAndroid = true
							}
						}
					}
				}
				if !disconnectedAndroid && !util.ContainsAny(b.Builder, hungSlavesIgnore) {
					// Hung buildslave?
					for _, s := range b.Steps {
						if s.Name == "steps" {
							continue
						}
						// If the step has been running for over an hour, it's probably hung.
						if s.Finished == 0 && time.Since(time.Unix(int64(s.Started), 0)) > hangTimePeriod {
							if err := am.AddAlert(&alerting.Alert{
								Name:        fmt.Sprintf("Possibly hung buildslave (%s)", b.BuildSlave),
								Category:    alerting.INFRA_ALERT,
								Message:     fmt.Sprintf(HUNG_BUILDSLAVE, b.BuildSlave, hangTimePeriod.String(), b.Master, b.Builder, b.Number, b.BuildSlave, b.BuildSlave),
								Nag:         int64(time.Hour),
								Actions:     actions,
								AutoDismiss: int64(10 * tickInterval),
							}); err != nil {
								glog.Error(err)
							}
						}
					}
				}
			}
		}
	}()

	// Failed update_scripts.
	go func() {
		lastSearch := time.Now()
		for _ = range time.Tick(tickInterval) {
			glog.Infof("Searching for builds which failed update_scripts.")
			currentSearch := time.Now()
			builds, err := buildbot.GetBuildsFromDateRange(lastSearch, currentSearch)
			lastSearch = currentSearch
			if err != nil {
				glog.Error(err)
				continue
			}
			for _, b := range builds {
				for _, s := range b.Steps {
					if s.Name == "update_scripts" {
						if s.Results != 0 {
							if err := am.AddAlert(&alerting.Alert{
								Name:     "update_scripts failed",
								Category: alerting.INFRA_ALERT,
								Message:  fmt.Sprintf(UPDATE_SCRIPTS, b.Builder, b.Master, b.Builder, b.Number, b.Builder, b.BuildSlave),
								Actions:  actions,
							}); err != nil {
								glog.Error(err)
							}
						}
						break
					}
				}
			}
		}
	}()
}
<|endoftext|>


<commit_before>package tarpan

import (
	"errors"
	"strconv"
	//"strings"
	//"runtime"
	"os"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	g "github.com/soniah/gosnmp"
)

// Tarpan constants
const (
	ConcurrentProcesses = 5
)

// Exit code
const (
	ExitCodeOK = iota
	ExitCodeConfError
	ExitCodeParseFlagError
)

// SNMP constants
const (
	COMMUNITY = "public"
	PORT      = 161
	TIMEOUT   = 2
	RETRY     = 3
	//MaximumPDUBytes = 512
)

// JSON structure definition
type DataSet struct {
	Default Defaults `json:"defaults"`
	Targets []Target `json:"targets"`
}

type Defaults struct {
	Port      uint16 `json:"port"`
	Version   string `json:"version"`
	Community string `json:"community"`
}

type Target struct {
	Name      string `json:"name"`
	Address   string `json:"address"`
	Port      uint16 `json:"port"`
	Version   string `json:"version"`
	Community string `json:"community"`
	OIDs      []OID  `json:"oids"`
}

type OID struct {
	OID         string `json:"oid"`
	Description string `json:"description"`
}

// Tarpan result structures
type TarpanResult struct {
	Name      string    `json:"name"`
	Address   string    `json:"address"`
	Port      uint16    `json:"port"`
	Version   string    `json:"version"`
	Community string    `json:"community"`
	VarBinds  []VarBind `json:"varbinds"`
}

type TarpanResults []*TarpanResult

type VarBind struct {
	Description string      `json:"description"`
	OID         string      `json:"oid"`
	Type        string      `json:"type"`
	Value       interface{} `json:"value"`
	Time        int64       `json:"time"`
}

type Tarpan interface {
	makeRequestBody()
	SetManager()
	SetParams()
	Get(params map[string]string, oids []string) ([]g.SnmpPDU, error)
	Run()
}

type TarpanManager struct {
	target *Target
	snmp   *g.GoSNMP
}

type TarpanManagers []*TarpanManager

type RequestParams struct {
	managerIndex int
	address      string
	community    string
	version      string
	port         string
	timeout      uint8
	retry        uint8
}

type Channels struct {
	semaphoe chan int
	results  chan TarpanResult
}

func (m *TarpanManager) setTarget(ds *DataSet, target_index int) error {
	params, param_err := getRequestParams(ds, target_index)
	if param_err != nil {
		log.Error(param_err)
		errors.New(param_err.Error())
	}
	m.SetParams(params)
	m.target = &ds.Targets[target_index]

	return nil
}

func (m *TarpanManager) Get(oids []string) (TarpanResult, error) {
	connection_err := m.snmp.Connect()
	defer m.snmp.Conn.Close()
	if connection_err != nil {
		log.Error("Connection error: %")
		return TarpanResult{}, errors.New(connection_err.Error())
	}
	snmpPacket, request_err := m.snmp.Get(oids)
	if request_err != nil {
		return TarpanResult{}, errors.New(request_err.Error())
	}
	tarpanResult := m.makeTarpanResult(snmpPacket)

	return tarpanResult, nil
}

func (m *TarpanManager) SetManager(manager *g.GoSNMP) {
	m.snmp = manager
}

func (m *TarpanManager) SetParams(p *RequestParams) {
	if p.address != "" {
		m.snmp.Target = p.address
	}
	if p.port != "" {
		port, _ := strconv.ParseUint(p.port, 10, 16)
		m.snmp.Port = uint16(port)
	} else {
		m.snmp.Port = uint16(PORT)
	}
	if p.community != "" {
		m.snmp.Community = p.community
	} else {
		m.snmp.Community = COMMUNITY
	}
	if p.version != "" {
		if p.version == "2c" {
			m.snmp.Version = g.Version2c
		}
	} else {
		m.snmp.Version = g.Version2c
	}
	if p.timeout != 0 {
		m.snmp.Timeout = time.Duration(p.timeout) * time.Second
	} else {
		m.snmp.Timeout = time.Duration(TIMEOUT) * time.Second
	}
	if p.retry != 0 {
		m.snmp.Retries = int(p.retry)
	} else {
		m.snmp.Retries = RETRY
	}
}

func getRequestParams(ds *DataSet, idx int) (*RequestParams, error) {
	var address string
	err := validateIP(ds.Targets[idx].Address)
	if err == nil {
		address = ds.Targets[idx].Address
	} else {
		return &RequestParams{}, err
	}
	community := ds.Targets[idx].Community
	if community == "" {
		community = ds.Default.Community
	}
	version := ds.Targets[idx].Version
	if version == "" {
		version = ds.Default.Version
	}
	port := ds.Targets[idx].Port
	if port == 0 {
		port = ds.Default.Port
	}

	requestParams := &RequestParams{
		managerIndex: idx,
		address:      address,
		community:    community,
		version:      version,
		port:         strconv.Itoa(int(port)),
	}

	return requestParams, nil
}

func makeChannel(result_buffer_size int) *Channels {
	channel := &Channels{
		semaphoe: make(chan int, ConcurrentProcesses),
		results:  make(chan TarpanResult, result_buffer_size),
	}

	return channel
}

func (m *TarpanManager) getTargetOIDDescription(oid string) (string, error) {
	for _, o := range m.target.OIDs {
		if o.OID == oid {
			return o.Description, nil
		}
	}

	return "", errors.New("oid not found")
}

func removeNilResult(tarpanResults []*TarpanResult) []*TarpanResult {
	var responseResult []*TarpanResult
	for _, v := range tarpanResults {
		if v != nil {
			responseResult = append(responseResult, v)
		}
	}

	return responseResult
}

func (m *TarpanManager) makeTarpanResult(sp *g.SnmpPacket) TarpanResult {
	now := time.Now()
	var varbinds []VarBind
	for _, val := range sp.Variables {
		//n := strings.TrimLeft(val.Name, ".")
		desc, _ := m.getTargetOIDDescription(val.Name)
		asn1ber_name, _ := getAsn1BERName(val.Type)
		//value := formatSnmpValue(val.Type, val.Value.(string))
		v := VarBind{
			Description: desc,
			OID:         val.Name,
			Type:        asn1ber_name,
			Value:       val.Value,
			Time:        now.Unix(),
		}
		varbinds = append(varbinds, v)
	}
	version := getSnmpVersionString(m.snmp.Version)
	tarpanResult := TarpanResult{
		Name:      m.target.Name,
		Address:   m.target.Address,
		Port:      m.snmp.Port,
		Version:   version,
		Community: m.snmp.Community,
		VarBinds:  varbinds,
	}

	return tarpanResult
}

func makeTarpanResults(c *Channels) []*TarpanResult {
	result_length := len(c.results)
	tarpanResults := make(TarpanResults, result_length)
	for i := 0; i < result_length; i++ {
		tarpanResult := <-c.results
		tarpanResults[i] = &tarpanResult
	}

	return removeNilResult(tarpanResults)
}

func makeManagers(ds *DataSet) []*TarpanManager {
	var manager *TarpanManager
	var managers TarpanManagers
	var oids []string

	// Set Target
	for t_idx := range ds.Targets {
		// TODO: split oid slice depending on PDU size
		oids = []string{}
		for o_idx := range ds.Targets[t_idx].OIDs {
			oids = append(oids, ds.Targets[t_idx].OIDs[o_idx].OID)
		}
		manager = &TarpanManager{
			snmp: &g.GoSNMP{},
		}
		err := manager.setTarget(ds, t_idx)
		if err != nil {
			log.Error("set target error")
		}
		managers = append(managers, manager)
	}

	return managers
}

func Collect(dataset *DataSet) []*TarpanResult {
	var waitGroup sync.WaitGroup
	var oids []string
	var tarpanResults []*TarpanResult

	managers := makeManagers(dataset)
	channels := makeChannel(len(dataset.Targets))
	for t_idx := range dataset.Targets {

		// TODO: split oid slice depending on PDU size
		oids = []string{}
		for o_idx := range dataset.Targets[t_idx].OIDs {
			oids = append(oids, dataset.Targets[t_idx].OIDs[o_idx].OID)
		}

		waitGroup.Add(1)
		go func(m *TarpanManager, o []string, c *Channels) {
			c.semaphoe <- 0
			defer func() {
				waitGroup.Done()
				<-c.semaphoe
			}()
			result, err := m.Get(o)
			if err != nil {
				log.Error(err)
				return
			}
			c.results <- result
		}(managers[t_idx], oids, channels)
	}
	waitGroup.Wait()
	tarpanResults = makeTarpanResults(channels)

	return tarpanResults
}

func Run(target string, output string, debug bool) (int, error) {
	var err error

	log.SetOutput(os.Stderr)
	if debug == true {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.WarnLevel)
	}

	// Load targets
	dataset, loadErr := loadConfig(target)
	if loadErr != nil {
		log.Error("err message:")
		return ExitCodeConfError, loadErr
	}

	// collect data
	tr := Collect(dataset)

	// output
	Write(output, tr)

	return ExitCodeOK, err
}
<commit_msg>change function order<commit_after>package tarpan

import (
	"errors"
	"os"
	"strconv"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	g "github.com/soniah/gosnmp"
)

// Tarpan constants
const (
	ConcurrentProcesses = 5
)

// Exit code
const (
	ExitCodeOK = iota
	ExitCodeConfError
	ExitCodeParseFlagError
)

// SNMP constants
const (
	COMMUNITY = "public"
	PORT      = 161
	TIMEOUT   = 2
	RETRY     = 3
	//MaximumPDUBytes = 512
)

// JSON structure definition
type DataSet struct {
	Default Defaults `json:"defaults"`
	Targets []Target `json:"targets"`
}

type Defaults struct {
	Port      uint16 `json:"port"`
	Version   string `json:"version"`
	Community string `json:"community"`
}

type Target struct {
	Name      string `json:"name"`
	Address   string `json:"address"`
	Port      uint16 `json:"port"`
	Version   string `json:"version"`
	Community string `json:"community"`
	OIDs      []OID  `json:"oids"`
}

type OID struct {
	OID         string `json:"oid"`
	Description string `json:"description"`
}

// Tarpan result structures
type TarpanResult struct {
	Name      string    `json:"name"`
	Address   string    `json:"address"`
	Port      uint16    `json:"port"`
	Version   string    `json:"version"`
	Community string    `json:"community"`
	VarBinds  []VarBind `json:"varbinds"`
}

type TarpanResults []*TarpanResult

type VarBind struct {
	Description string      `json:"description"`
	OID         string      `json:"oid"`
	Type        string      `json:"type"`
	Value       interface{} `json:"value"`
	Time        int64       `json:"time"`
}

type Tarpan interface {
	setTarget(ds *DataSet, target_index int)
	Get(oids []string)
	SetManager()
	setParams(p *RequestParams)
	getTargetOIDDescription(oid string)
	makeTarpanResult(sp *g.SnmpPacket)
	Run()
}

type TarpanManager struct {
	target *Target
	snmp   *g.GoSNMP
}

type TarpanManagers []*TarpanManager

type RequestParams struct {
	managerIndex int
	address      string
	community    string
	version      string
	port         string
	timeout      uint8
	retry        uint8
}

type Channels struct {
	semaphoe chan int
	results  chan TarpanResult
}

func (m *TarpanManager) setTarget(ds *DataSet, target_index int) error {
	params, param_err := getRequestParams(ds, target_index)
	if param_err != nil {
		log.Error(param_err)
		errors.New(param_err.Error())
	}
	m.setParams(params)
	m.target = &ds.Targets[target_index]

	return nil
}

func (m *TarpanManager) Get(oids []string) (TarpanResult, error) {
	connection_err := m.snmp.Connect()
	defer m.snmp.Conn.Close()
	if connection_err != nil {
		log.Error("Connection error: %")
		return TarpanResult{}, errors.New(connection_err.Error())
	}
	snmpPacket, request_err := m.snmp.Get(oids)
	if request_err != nil {
		return TarpanResult{}, errors.New(request_err.Error())
	}
	tarpanResult := m.makeTarpanResult(snmpPacket)

	return tarpanResult, nil
}

func (m *TarpanManager) SetManager(manager *g.GoSNMP) {
	m.snmp = manager
}

func (m *TarpanManager) setParams(p *RequestParams) {
	if p.address != "" {
		m.snmp.Target = p.address
	}
	if p.port != "" {
		port, _ := strconv.ParseUint(p.port, 10, 16)
		m.snmp.Port = uint16(port)
	} else {
		m.snmp.Port = uint16(PORT)
	}
	if p.community != "" {
		m.snmp.Community = p.community
	} else {
		m.snmp.Community = COMMUNITY
	}
	if p.version != "" {
		if p.version == "2c" {
			m.snmp.Version = g.Version2c
		}
	} else {
		m.snmp.Version = g.Version2c
	}
	if p.timeout != 0 {
		m.snmp.Timeout = time.Duration(p.timeout) * time.Second
	} else {
		m.snmp.Timeout = time.Duration(TIMEOUT) * time.Second
	}
	if p.retry != 0 {
		m.snmp.Retries = int(p.retry)
	} else {
		m.snmp.Retries = RETRY
	}
}

func getRequestParams(ds *DataSet, idx int) (*RequestParams, error) {
	var address string
	err := validateIP(ds.Targets[idx].Address)
	if err == nil {
		address = ds.Targets[idx].Address
	} else {
		return &RequestParams{}, err
	}
	community := ds.Targets[idx].Community
	if community == "" {
		community = ds.Default.Community
	}
	version := ds.Targets[idx].Version
	if version == "" {
		version = ds.Default.Version
	}
	port := ds.Targets[idx].Port
	if port == 0 {
		port = ds.Default.Port
	}

	requestParams := &RequestParams{
		managerIndex: idx,
		address:      address,
		community:    community,
		version:      version,
		port:         strconv.Itoa(int(port)),
	}

	return requestParams, nil
}

func makeChannel(result_buffer_size int) *Channels {
	channel := &Channels{
		semaphoe: make(chan int, ConcurrentProcesses),
		results:  make(chan TarpanResult, result_buffer_size),
	}

	return channel
}

func (m *TarpanManager) getTargetOIDDescription(oid string) (string, error) {
	for _, o := range m.target.OIDs {
		if o.OID == oid {
			return o.Description, nil
		}
	}

	return "", errors.New("oid not found")
}

func removeNilResult(tarpanResults []*TarpanResult) []*TarpanResult {
	var responseResult []*TarpanResult
	for _, v := range tarpanResults {
		if v != nil {
			responseResult = append(responseResult, v)
		}
	}

	return responseResult
}

func (m *TarpanManager) makeTarpanResult(sp *g.SnmpPacket) TarpanResult {
	now := time.Now()
	var varbinds []VarBind
	for _, val := range sp.Variables {
		//n := strings.TrimLeft(val.Name, ".")
		desc, _ := m.getTargetOIDDescription(val.Name)
		asn1ber_name, _ := getAsn1BERName(val.Type)
		//value := formatSnmpValue(val.Type, val.Value.(string))
		v := VarBind{
			Description: desc,
			OID:         val.Name,
			Type:        asn1ber_name,
			Value:       val.Value,
			Time:        now.Unix(),
		}
		varbinds = append(varbinds, v)
	}
	version := getSnmpVersionString(m.snmp.Version)
	tarpanResult := TarpanResult{
		Name:      m.target.Name,
		Address:   m.target.Address,
		Port:      m.snmp.Port,
		Version:   version,
		Community: m.snmp.Community,
		VarBinds:  varbinds,
	}

	return tarpanResult
}

func makeTarpanResults(c *Channels) []*TarpanResult {
	result_length := len(c.results)
	tarpanResults := make(TarpanResults, result_length)
	for i := 0; i < result_length; i++ {
		tarpanResult := <-c.results
		tarpanResults[i] = &tarpanResult
	}

	return removeNilResult(tarpanResults)
}

func makeManagers(ds *DataSet) []*TarpanManager {
	var manager *TarpanManager
	var managers TarpanManagers

	// Set Target
	for t_idx := range ds.Targets {
		manager = &TarpanManager{
			snmp: &g.GoSNMP{},
		}
		err := manager.setTarget(ds, t_idx)
		if err != nil {
			log.Error("set target error")
		}
		managers = append(managers, manager)
	}

	return managers
}

func Collect(dataset *DataSet) []*TarpanResult {
	var waitGroup sync.WaitGroup
	var oids []string
	var tarpanResults []*TarpanResult

	managers := makeManagers(dataset)
	channels := makeChannel(len(dataset.Targets))
	for t_idx := range dataset.Targets {

		// TODO: split oid slice depending on PDU size
		oids = []string{}
		for o_idx := range dataset.Targets[t_idx].OIDs {
			oids = append(oids, dataset.Targets[t_idx].OIDs[o_idx].OID)
		}

		waitGroup.Add(1)
		go func(m *TarpanManager, o []string, c *Channels) {
			c.semaphoe <- 0
			defer func() {
				waitGroup.Done()
				<-c.semaphoe
			}()
			result, err := m.Get(o)
			if err != nil {
				log.Error(err)
				return
			}
			c.results <- result
		}(managers[t_idx], oids, channels)
	}
	waitGroup.Wait()
	tarpanResults = makeTarpanResults(channels)

	return tarpanResults
}

func Run(target string, output_type string, debug bool) (int, error) {
	var err error

	// setup logger
	log.SetOutput(os.Stderr)
	if debug == true {
		log.SetLevel(log.DebugLevel)
	} else {
		log.SetLevel(log.WarnLevel)
	}

	// load targets
	dataset, loadErr := loadConfig(target)
	if loadErr != nil {
		log.Error(loadErr)
		return ExitCodeConfError, loadErr
	}

	// collect data
	tr := Collect(dataset)

	// output
	Write(output_type, tr)

	return ExitCodeOK, err
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess Authors.

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

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

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

package schema

import (
	"context"
	"sync"
	"time"

	"vitess.io/vitess/go/vt/callerid"

	"vitess.io/vitess/go/vt/vttablet/queryservice"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	querypb "vitess.io/vitess/go/vt/proto/query"

	"vitess.io/vitess/go/vt/discovery"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vtgate/vindexes"
)

type (
	keyspaceStr  = string
	tableNameStr = string

	// Tracker contains the required fields to perform schema tracking.
	Tracker struct {
		ch     chan *discovery.TabletHealth
		cancel context.CancelFunc

		mu     sync.Mutex
		tables *tableMap
		ctx    context.Context
		signal func() // a function that we'll call whenever we have new schema data

		// map of keyspace currently tracked
		tracked      map[keyspaceStr]*updateController
		consumeDelay time.Duration
	}
)

// defaultConsumeDelay is the default time, the updateController will wait before checking the schema fetch request queue.
const defaultConsumeDelay = 1 * time.Second

// NewTracker creates the tracker object.
func NewTracker(ch chan *discovery.TabletHealth, user *string) *Tracker {
	ctx := context.Background()
	// Set the caller on the context if the user is provided.
	// This user that will be sent down to vttablet calls.
	if user != nil && *user != "" {
		ctx = callerid.NewContext(ctx, nil, callerid.NewImmediateCallerID(*user))
	}

	return &Tracker{
		ctx:          ctx,
		ch:           ch,
		tables:       &tableMap{m: map[keyspaceStr]map[tableNameStr][]vindexes.Column{}},
		tracked:      map[keyspaceStr]*updateController{},
		consumeDelay: defaultConsumeDelay,
	}
}

// LoadKeyspace loads the keyspace schema.
func (t *Tracker) LoadKeyspace(conn queryservice.QueryService, target *querypb.Target) error {
	res, err := conn.Execute(t.ctx, target, mysql.FetchTables, nil, 0, 0, nil)
	if err != nil {
		return err
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	// We must clear out any previous schema before loading it here as this is called
	// whenever a shard's primary tablet starts and sends the initial signal. Without
	// clearing out the previous schema we can end up with duplicate entries when the
	// tablet is simply restarted or potentially when we elect a new primary.
	t.clearKeyspaceTables(target.Keyspace)
	t.updateTables(target.Keyspace, res)
	t.tracked[target.Keyspace].setLoaded(true)
	log.Infof("finished loading schema for keyspace %s. Found %d columns in total across the tables", target.Keyspace, len(res.Rows))
	return nil
}

// Start starts the schema tracking.
func (t *Tracker) Start() {
	log.Info("Starting schema tracking")
	ctx, cancel := context.WithCancel(t.ctx)
	t.cancel = cancel
	go func(ctx context.Context, t *Tracker) {
		for {
			select {
			case th := <-t.ch:
				ksUpdater := t.getKeyspaceUpdateController(th)
				ksUpdater.add(th)
			case <-ctx.Done():
				// closing of the channel happens outside the scope of the tracker. It is the responsibility of the one who created this tracker.
				return
			}
		}
	}(ctx, t)
}

// getKeyspaceUpdateController returns the updateController for the given keyspace
// the updateController will be created if there was none.
func (t *Tracker) getKeyspaceUpdateController(th *discovery.TabletHealth) *updateController {
	t.mu.Lock()
	defer t.mu.Unlock()

	ksUpdater, exists := t.tracked[th.Target.Keyspace]
	if !exists {
		ksUpdater = t.newUpdateController()
		t.tracked[th.Target.Keyspace] = ksUpdater
	}
	return ksUpdater
}

func (t *Tracker) newUpdateController() *updateController {
	return &updateController{update: t.updateSchema, reloadKeyspace: t.initKeyspace, signal: t.signal, consumeDelay: t.consumeDelay}
}

func (t *Tracker) initKeyspace(th *discovery.TabletHealth) error {
	err := t.LoadKeyspace(th.Conn, th.Target)
	if err != nil {
		log.Warningf("Unable to add keyspace to tracker: %v", err)
		return err
	}
	return nil
}

// Stop stops the schema tracking
func (t *Tracker) Stop() {
	log.Info("Stopping schema tracking")
	t.cancel()
}

// GetColumns returns the column list for table in the given keyspace.
func (t *Tracker) GetColumns(ks string, tbl string) []vindexes.Column {
	t.mu.Lock()
	defer t.mu.Unlock()

	return t.tables.get(ks, tbl)
}

// Tables returns a map with the columns for all known tables in the keyspace
func (t *Tracker) Tables(ks string) map[string][]vindexes.Column {
	t.mu.Lock()
	defer t.mu.Unlock()

	m := t.tables.m[ks]
	if m == nil {
		return map[string][]vindexes.Column{} // we know nothing about this KS, so that is the info we can give out
	}

	return m
}

func (t *Tracker) updateSchema(th *discovery.TabletHealth) bool {
	tablesUpdated := th.Stats.TableSchemaChanged
	tables, err := sqltypes.BuildBindVariable(tablesUpdated)
	if err != nil {
		log.Errorf("failed to read updated tables from TabletHealth: %v", err)
		return false
	}
	bv := map[string]*querypb.BindVariable{"tableNames": tables}
	res, err := th.Conn.Execute(t.ctx, th.Target, mysql.FetchUpdatedTables, bv, 0, 0, nil)
	if err != nil {
		t.tracked[th.Target.Keyspace].setLoaded(false)
		// TODO: optimize for the tables that got errored out.
		log.Warningf("error fetching new schema for %v, making them non-authoritative: %v", tablesUpdated, err)
		return false
	}

	t.mu.Lock()
	defer t.mu.Unlock()

	// first we empty all prior schema. deleted tables will not show up in the result,
	// so this is the only chance to delete
	for _, tbl := range tablesUpdated {
		t.tables.delete(th.Target.Keyspace, tbl)
	}
	t.updateTables(th.Target.Keyspace, res)
	return true
}

func (t *Tracker) updateTables(keyspace string, res *sqltypes.Result) {
	for _, row := range res.Rows {
		tbl := row[0].ToString()
		colName := row[1].ToString()
		colType := row[2].ToString()
		collation := row[3].ToString()

		cType := sqlparser.ColumnType{Type: colType}
		col := vindexes.Column{Name: sqlparser.NewColIdent(colName), Type: cType.SQLType(), CollationName: collation}
		cols := t.tables.get(keyspace, tbl)

		t.tables.set(keyspace, tbl, append(cols, col))
	}
}

// RegisterSignalReceiver allows a function to register to be called when new schema is available
func (t *Tracker) RegisterSignalReceiver(f func()) {
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, controller := range t.tracked {
		controller.signal = f
	}
	t.signal = f
}

// AddNewKeyspace adds keyspace to the tracker.
func (t *Tracker) AddNewKeyspace(conn queryservice.QueryService, target *querypb.Target) error {
	updateController := t.newUpdateController()
	t.tracked[target.Keyspace] = updateController
	err := t.LoadKeyspace(conn, target)
	if err != nil {
		updateController.setIgnore(checkIfWeShouldIgnoreKeyspace(err))
	}
	return err
}

type tableMap struct {
	m map[keyspaceStr]map[tableNameStr][]vindexes.Column
}

func (tm *tableMap) set(ks, tbl string, cols []vindexes.Column) {
	m := tm.m[ks]
	if m == nil {
		m = make(map[tableNameStr][]vindexes.Column)
		tm.m[ks] = m
	}
	m[tbl] = cols
}

func (tm *tableMap) get(ks, tbl string) []vindexes.Column {
	m := tm.m[ks]
	if m == nil {
		return nil
	}
	return m[tbl]
}

func (tm *tableMap) delete(ks, tbl string) {
	m := tm.m[ks]
	if m == nil {
		return
	}
	delete(m, tbl)
}

// This empties out any previous schema for for all tables in a keyspace.
// You should call this before initializing/loading a keyspace of the same
// name in the cache.
func (t *Tracker) clearKeyspaceTables(ks string) {
	if t.tables != nil && t.tables.m != nil {
		delete(t.tables.m, ks)
	}
}
<commit_msg>Improve initKeyspace failure logging for easier debugging<commit_after>/*
Copyright 2021 The Vitess Authors.

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

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

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

package schema

import (
	"context"
	"sync"
	"time"

	"vitess.io/vitess/go/vt/callerid"

	"vitess.io/vitess/go/vt/vttablet/queryservice"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sqltypes"
	querypb "vitess.io/vitess/go/vt/proto/query"

	"vitess.io/vitess/go/vt/discovery"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/sqlparser"
	"vitess.io/vitess/go/vt/vtgate/vindexes"
)

type (
	keyspaceStr  = string
	tableNameStr = string

	// Tracker contains the required fields to perform schema tracking.
	Tracker struct {
		ch     chan *discovery.TabletHealth
		cancel context.CancelFunc

		mu     sync.Mutex
		tables *tableMap
		ctx    context.Context
		signal func() // a function that we'll call whenever we have new schema data

		// map of keyspace currently tracked
		tracked      map[keyspaceStr]*updateController
		consumeDelay time.Duration
	}
)

// defaultConsumeDelay is the default time, the updateController will wait before checking the schema fetch request queue.
const defaultConsumeDelay = 1 * time.Second

// NewTracker creates the tracker object.
func NewTracker(ch chan *discovery.TabletHealth, user *string) *Tracker {
	ctx := context.Background()
	// Set the caller on the context if the user is provided.
	// This user that will be sent down to vttablet calls.
	if user != nil && *user != "" {
		ctx = callerid.NewContext(ctx, nil, callerid.NewImmediateCallerID(*user))
	}

	return &Tracker{
		ctx:          ctx,
		ch:           ch,
		tables:       &tableMap{m: map[keyspaceStr]map[tableNameStr][]vindexes.Column{}},
		tracked:      map[keyspaceStr]*updateController{},
		consumeDelay: defaultConsumeDelay,
	}
}

// LoadKeyspace loads the keyspace schema.
func (t *Tracker) LoadKeyspace(conn queryservice.QueryService, target *querypb.Target) error {
	res, err := conn.Execute(t.ctx, target, mysql.FetchTables, nil, 0, 0, nil)
	if err != nil {
		return err
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	// We must clear out any previous schema before loading it here as this is called
	// whenever a shard's primary tablet starts and sends the initial signal. Without
	// clearing out the previous schema we can end up with duplicate entries when the
	// tablet is simply restarted or potentially when we elect a new primary.
	t.clearKeyspaceTables(target.Keyspace)
	t.updateTables(target.Keyspace, res)
	t.tracked[target.Keyspace].setLoaded(true)
	log.Infof("finished loading schema for keyspace %s. Found %d columns in total across the tables", target.Keyspace, len(res.Rows))
	return nil
}

// Start starts the schema tracking.
func (t *Tracker) Start() {
	log.Info("Starting schema tracking")
	ctx, cancel := context.WithCancel(t.ctx)
	t.cancel = cancel
	go func(ctx context.Context, t *Tracker) {
		for {
			select {
			case th := <-t.ch:
				ksUpdater := t.getKeyspaceUpdateController(th)
				ksUpdater.add(th)
			case <-ctx.Done():
				// closing of the channel happens outside the scope of the tracker. It is the responsibility of the one who created this tracker.
				return
			}
		}
	}(ctx, t)
}

// getKeyspaceUpdateController returns the updateController for the given keyspace
// the updateController will be created if there was none.
func (t *Tracker) getKeyspaceUpdateController(th *discovery.TabletHealth) *updateController {
	t.mu.Lock()
	defer t.mu.Unlock()

	ksUpdater, exists := t.tracked[th.Target.Keyspace]
	if !exists {
		ksUpdater = t.newUpdateController()
		t.tracked[th.Target.Keyspace] = ksUpdater
	}
	return ksUpdater
}

func (t *Tracker) newUpdateController() *updateController {
	return &updateController{update: t.updateSchema, reloadKeyspace: t.initKeyspace, signal: t.signal, consumeDelay: t.consumeDelay}
}

func (t *Tracker) initKeyspace(th *discovery.TabletHealth) error {
	err := t.LoadKeyspace(th.Conn, th.Target)
	if err != nil {
		log.Warningf("Unable to add the %s keyspace to the schema tracker: %v", th.Target.Keyspace, err)
		return err
	}
	return nil
}

// Stop stops the schema tracking
func (t *Tracker) Stop() {
	log.Info("Stopping schema tracking")
	t.cancel()
}

// GetColumns returns the column list for table in the given keyspace.
func (t *Tracker) GetColumns(ks string, tbl string) []vindexes.Column {
	t.mu.Lock()
	defer t.mu.Unlock()

	return t.tables.get(ks, tbl)
}

// Tables returns a map with the columns for all known tables in the keyspace
func (t *Tracker) Tables(ks string) map[string][]vindexes.Column {
	t.mu.Lock()
	defer t.mu.Unlock()

	m := t.tables.m[ks]
	if m == nil {
		return map[string][]vindexes.Column{} // we know nothing about this KS, so that is the info we can give out
	}

	return m
}

func (t *Tracker) updateSchema(th *discovery.TabletHealth) bool {
	tablesUpdated := th.Stats.TableSchemaChanged
	tables, err := sqltypes.BuildBindVariable(tablesUpdated)
	if err != nil {
		log.Errorf("failed to read updated tables from TabletHealth: %v", err)
		return false
	}
	bv := map[string]*querypb.BindVariable{"tableNames": tables}
	res, err := th.Conn.Execute(t.ctx, th.Target, mysql.FetchUpdatedTables, bv, 0, 0, nil)
	if err != nil {
		t.tracked[th.Target.Keyspace].setLoaded(false)
		// TODO: optimize for the tables that got errored out.
		log.Warningf("error fetching new schema for %v, making them non-authoritative: %v", tablesUpdated, err)
		return false
	}

	t.mu.Lock()
	defer t.mu.Unlock()

	// first we empty all prior schema. deleted tables will not show up in the result,
	// so this is the only chance to delete
	for _, tbl := range tablesUpdated {
		t.tables.delete(th.Target.Keyspace, tbl)
	}
	t.updateTables(th.Target.Keyspace, res)
	return true
}

func (t *Tracker) updateTables(keyspace string, res *sqltypes.Result) {
	for _, row := range res.Rows {
		tbl := row[0].ToString()
		colName := row[1].ToString()
		colType := row[2].ToString()
		collation := row[3].ToString()

		cType := sqlparser.ColumnType{Type: colType}
		col := vindexes.Column{Name: sqlparser.NewColIdent(colName), Type: cType.SQLType(), CollationName: collation}
		cols := t.tables.get(keyspace, tbl)

		t.tables.set(keyspace, tbl, append(cols, col))
	}
}

// RegisterSignalReceiver allows a function to register to be called when new schema is available
func (t *Tracker) RegisterSignalReceiver(f func()) {
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, controller := range t.tracked {
		controller.signal = f
	}
	t.signal = f
}

// AddNewKeyspace adds keyspace to the tracker.
func (t *Tracker) AddNewKeyspace(conn queryservice.QueryService, target *querypb.Target) error {
	updateController := t.newUpdateController()
	t.tracked[target.Keyspace] = updateController
	err := t.LoadKeyspace(conn, target)
	if err != nil {
		updateController.setIgnore(checkIfWeShouldIgnoreKeyspace(err))
	}
	return err
}

type tableMap struct {
	m map[keyspaceStr]map[tableNameStr][]vindexes.Column
}

func (tm *tableMap) set(ks, tbl string, cols []vindexes.Column) {
	m := tm.m[ks]
	if m == nil {
		m = make(map[tableNameStr][]vindexes.Column)
		tm.m[ks] = m
	}
	m[tbl] = cols
}

func (tm *tableMap) get(ks, tbl string) []vindexes.Column {
	m := tm.m[ks]
	if m == nil {
		return nil
	}
	return m[tbl]
}

func (tm *tableMap) delete(ks, tbl string) {
	m := tm.m[ks]
	if m == nil {
		return
	}
	delete(m, tbl)
}

// This empties out any previous schema for for all tables in a keyspace.
// You should call this before initializing/loading a keyspace of the same
// name in the cache.
func (t *Tracker) clearKeyspaceTables(ks string) {
	if t.tables != nil && t.tables.m != nil {
		delete(t.tables.m, ks)
	}
}
<|endoftext|>


<commit_before>package models

import (
	"time"
)

// SubscriptionInfo shows how many times user is subscribed to a channel
type SubscriptionInfo struct {
	Count     int       `json:"count"`
	IsPrime   bool      `json:"isPrime"`
	User      string    `json:"user"`
	UserID    string    `json:"userID"`
	ChannelID string    `json:"channelID"`
	Date      time.Time `json:"date"`
}
<commit_msg>Feature: Added id field to subscription info for bookmarking<commit_after>package models

import (
	"time"

	"gopkg.in/mgo.v2/bson"
)

// SubscriptionInfo shows how many t imes user is subscribed to a channel
type SubscriptionInfo struct {
	ID        bson.ObjectId `bson:"_id,omitempty" json:"id"`
	Count     int           `json:"count"`
	IsPrime   bool          `json:"isPrime"`
	User      string        `json:"user"`
	UserID    string        `json:"userID"`
	ChannelID string        `json:"channelID"`
	Date      time.Time     `json:"date"`
}
<|endoftext|>


<commit_before>package pool

import (
	"context"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-redis/redis/v8/internal"
)

var ErrClosed = errors.New("redis: client is closed")
var ErrPoolTimeout = errors.New("redis: connection pool timeout")

var timers = sync.Pool{
	New: func() interface{} {
		t := time.NewTimer(time.Hour)
		t.Stop()
		return t
	},
}

// Stats contains pool state information and accumulated stats.
type Stats struct {
	Hits     uint32 // number of times free connection was found in the pool
	Misses   uint32 // number of times free connection was NOT found in the pool
	Timeouts uint32 // number of times a wait timeout occurred

	TotalConns uint32 // number of total connections in the pool
	IdleConns  uint32 // number of idle connections in the pool
	StaleConns uint32 // number of stale connections removed from the pool
}

type Pooler interface {
	NewConn(context.Context) (*Conn, error)
	CloseConn(*Conn) error

	Get(context.Context) (*Conn, error)
	Put(*Conn)
	Remove(*Conn, error)

	Len() int
	IdleLen() int
	Stats() *Stats

	Close() error
}

type Options struct {
	Dialer  func(context.Context) (net.Conn, error)
	OnClose func(*Conn) error

	PoolSize           int
	MinIdleConns       int
	MaxConnAge         time.Duration
	PoolTimeout        time.Duration
	IdleTimeout        time.Duration
	IdleCheckFrequency time.Duration
}

type ConnPool struct {
	opt *Options

	dialErrorsNum uint32 // atomic

	lastDialErrorMu sync.RWMutex
	lastDialError   error

	queue chan struct{}

	connsMu      sync.Mutex
	conns        []*Conn
	idleConns    []*Conn
	poolSize     int
	idleConnsLen int

	stats Stats

	_closed  uint32 // atomic
	closedCh chan struct{}
}

var _ Pooler = (*ConnPool)(nil)

func NewConnPool(opt *Options) *ConnPool {
	p := &ConnPool{
		opt: opt,

		queue:     make(chan struct{}, opt.PoolSize),
		conns:     make([]*Conn, 0, opt.PoolSize),
		idleConns: make([]*Conn, 0, opt.PoolSize),
		closedCh:  make(chan struct{}),
	}

	p.connsMu.Lock()
	p.checkMinIdleConns()
	p.connsMu.Unlock()

	if opt.IdleTimeout > 0 && opt.IdleCheckFrequency > 0 {
		go p.reaper(opt.IdleCheckFrequency)
	}

	return p
}

func (p *ConnPool) checkMinIdleConns() {
	if p.opt.MinIdleConns == 0 {
		return
	}
	for p.poolSize < p.opt.PoolSize && p.idleConnsLen < p.opt.MinIdleConns {
		p.poolSize++
		p.idleConnsLen++
		go func() {
			err := p.addIdleConn()
			if err != nil {
				p.connsMu.Lock()
				p.poolSize--
				p.idleConnsLen--
				p.connsMu.Unlock()
			}
		}()
	}
}

func (p *ConnPool) addIdleConn() error {
	cn, err := p.dialConn(context.TODO(), true)
	if err != nil {
		return err
	}

	p.connsMu.Lock()
	p.conns = append(p.conns, cn)
	p.idleConns = append(p.idleConns, cn)
	p.connsMu.Unlock()
	return nil
}

func (p *ConnPool) NewConn(ctx context.Context) (*Conn, error) {
	return p.newConn(ctx, false)
}

func (p *ConnPool) newConn(ctx context.Context, pooled bool) (*Conn, error) {
	cn, err := p.dialConn(ctx, pooled)
	if err != nil {
		return nil, err
	}

	p.connsMu.Lock()
	p.conns = append(p.conns, cn)
	if pooled {
		// If pool is full remove the cn on next Put.
		if p.poolSize >= p.opt.PoolSize {
			cn.pooled = false
		} else {
			p.poolSize++
		}
	}
	p.connsMu.Unlock()
	return cn, nil
}

func (p *ConnPool) dialConn(ctx context.Context, pooled bool) (*Conn, error) {
	if p.closed() {
		return nil, ErrClosed
	}

	if atomic.LoadUint32(&p.dialErrorsNum) >= uint32(p.opt.PoolSize) {
		return nil, p.getLastDialError()
	}

	netConn, err := p.opt.Dialer(ctx)
	if err != nil {
		p.setLastDialError(err)
		if atomic.AddUint32(&p.dialErrorsNum, 1) == uint32(p.opt.PoolSize) {
			go p.tryDial()
		}
		return nil, err
	}

	cn := NewConn(netConn)
	cn.pooled = pooled
	return cn, nil
}

func (p *ConnPool) tryDial() {
	for {
		if p.closed() {
			return
		}

		conn, err := p.opt.Dialer(context.Background())
		if err != nil {
			p.setLastDialError(err)
			time.Sleep(time.Second)
			continue
		}

		atomic.StoreUint32(&p.dialErrorsNum, 0)
		_ = conn.Close()
		return
	}
}

func (p *ConnPool) setLastDialError(err error) {
	p.lastDialErrorMu.Lock()
	p.lastDialError = err
	p.lastDialErrorMu.Unlock()
}

func (p *ConnPool) getLastDialError() error {
	p.lastDialErrorMu.RLock()
	err := p.lastDialError
	p.lastDialErrorMu.RUnlock()
	return err
}

// Get returns existed connection from the pool or creates a new one.
func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
	if p.closed() {
		return nil, ErrClosed
	}

	err := p.waitTurn(ctx)
	if err != nil {
		return nil, err
	}

	for {
		p.connsMu.Lock()
		cn := p.popIdle()
		p.connsMu.Unlock()

		if cn == nil {
			break
		}

		if p.isStaleConn(cn) {
			_ = p.CloseConn(cn)
			continue
		}

		atomic.AddUint32(&p.stats.Hits, 1)
		return cn, nil
	}

	atomic.AddUint32(&p.stats.Misses, 1)

	newcn, err := p.newConn(ctx, true)
	if err != nil {
		p.freeTurn()
		return nil, err
	}

	return newcn, nil
}

func (p *ConnPool) getTurn() {
	p.queue <- struct{}{}
}

func (p *ConnPool) waitTurn(ctx context.Context) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	select {
	case p.queue <- struct{}{}:
		return nil
	default:
	}

	timer := timers.Get().(*time.Timer)
	timer.Reset(p.opt.PoolTimeout)

	select {
	case <-ctx.Done():
		if !timer.Stop() {
			<-timer.C
		}
		timers.Put(timer)
		return ctx.Err()
	case p.queue <- struct{}{}:
		if !timer.Stop() {
			<-timer.C
		}
		timers.Put(timer)
		return nil
	case <-timer.C:
		timers.Put(timer)
		atomic.AddUint32(&p.stats.Timeouts, 1)
		return ErrPoolTimeout
	}
}

func (p *ConnPool) freeTurn() {
	<-p.queue
}

func (p *ConnPool) popIdle() *Conn {
	if len(p.idleConns) == 0 {
		return nil
	}

	idx := len(p.idleConns) - 1
	cn := p.idleConns[idx]
	p.idleConns = p.idleConns[:idx]
	p.idleConnsLen--
	p.checkMinIdleConns()
	return cn
}

func (p *ConnPool) Put(cn *Conn) {
	if cn.rd.Buffered() > 0 {
		internal.Logger.Printf("Conn has unread data")
		p.Remove(cn, BadConnError{})
		return
	}

	if !cn.pooled {
		p.Remove(cn, nil)
		return
	}

	p.connsMu.Lock()
	p.idleConns = append(p.idleConns, cn)
	p.idleConnsLen++
	p.connsMu.Unlock()
	p.freeTurn()
}

func (p *ConnPool) Remove(cn *Conn, reason error) {
	p.removeConnWithLock(cn)
	p.freeTurn()
	_ = p.closeConn(cn)
}

func (p *ConnPool) CloseConn(cn *Conn) error {
	p.removeConnWithLock(cn)
	return p.closeConn(cn)
}

func (p *ConnPool) removeConnWithLock(cn *Conn) {
	p.connsMu.Lock()
	p.removeConn(cn)
	p.connsMu.Unlock()
}

func (p *ConnPool) removeConn(cn *Conn) {
	for i, c := range p.conns {
		if c == cn {
			p.conns = append(p.conns[:i], p.conns[i+1:]...)
			if cn.pooled {
				p.poolSize--
				p.checkMinIdleConns()
			}
			return
		}
	}
}

func (p *ConnPool) closeConn(cn *Conn) error {
	if p.opt.OnClose != nil {
		_ = p.opt.OnClose(cn)
	}
	return cn.Close()
}

// Len returns total number of connections.
func (p *ConnPool) Len() int {
	p.connsMu.Lock()
	n := len(p.conns)
	p.connsMu.Unlock()
	return n
}

// IdleLen returns number of idle connections.
func (p *ConnPool) IdleLen() int {
	p.connsMu.Lock()
	n := p.idleConnsLen
	p.connsMu.Unlock()
	return n
}

func (p *ConnPool) Stats() *Stats {
	idleLen := p.IdleLen()
	return &Stats{
		Hits:     atomic.LoadUint32(&p.stats.Hits),
		Misses:   atomic.LoadUint32(&p.stats.Misses),
		Timeouts: atomic.LoadUint32(&p.stats.Timeouts),

		TotalConns: uint32(p.Len()),
		IdleConns:  uint32(idleLen),
		StaleConns: atomic.LoadUint32(&p.stats.StaleConns),
	}
}

func (p *ConnPool) closed() bool {
	return atomic.LoadUint32(&p._closed) == 1
}

func (p *ConnPool) Filter(fn func(*Conn) bool) error {
	var firstErr error
	p.connsMu.Lock()
	for _, cn := range p.conns {
		if fn(cn) {
			if err := p.closeConn(cn); err != nil && firstErr == nil {
				firstErr = err
			}
		}
	}
	p.connsMu.Unlock()
	return firstErr
}

func (p *ConnPool) Close() error {
	if !atomic.CompareAndSwapUint32(&p._closed, 0, 1) {
		return ErrClosed
	}
	close(p.closedCh)

	var firstErr error
	p.connsMu.Lock()
	for _, cn := range p.conns {
		if err := p.closeConn(cn); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	p.conns = nil
	p.poolSize = 0
	p.idleConns = nil
	p.idleConnsLen = 0
	p.connsMu.Unlock()

	return firstErr
}

func (p *ConnPool) reaper(frequency time.Duration) {
	ticker := time.NewTicker(frequency)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			// It is possible that ticker and closedCh arrive together,
			// and select pseudo-randomly pick ticker case, we double
			// check here to prevent being executed after closed.
			if p.closed() {
				return
			}
			_, err := p.ReapStaleConns()
			if err != nil {
				internal.Logger.Printf("ReapStaleConns failed: %s", err)
				continue
			}
		case <-p.closedCh:
			return
		}
	}
}

func (p *ConnPool) ReapStaleConns() (int, error) {
	var n int
	for {
		p.getTurn()

		p.connsMu.Lock()
		cn := p.reapStaleConn()
		p.connsMu.Unlock()
		p.freeTurn()

		if cn != nil {
			_ = p.closeConn(cn)
			n++
		} else {
			break
		}
	}
	atomic.AddUint32(&p.stats.StaleConns, uint32(n))
	return n, nil
}

func (p *ConnPool) reapStaleConn() *Conn {
	if len(p.idleConns) == 0 {
		return nil
	}

	cn := p.idleConns[0]
	if !p.isStaleConn(cn) {
		return nil
	}

	p.idleConns = append(p.idleConns[:0], p.idleConns[1:]...)
	p.idleConnsLen--
	p.removeConn(cn)

	return cn
}

func (p *ConnPool) isStaleConn(cn *Conn) bool {
	if p.opt.IdleTimeout == 0 && p.opt.MaxConnAge == 0 {
		return false
	}

	now := time.Now()
	if p.opt.IdleTimeout > 0 && now.Sub(cn.UsedAt()) >= p.opt.IdleTimeout {
		return true
	}
	if p.opt.MaxConnAge > 0 && now.Sub(cn.createdAt) >= p.opt.MaxConnAge {
		return true
	}

	return false
}
<commit_msg>Use atomic.Value instead of lock for ConnPool.lastDialError<commit_after>package pool

import (
	"context"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-redis/redis/v8/internal"
)

var ErrClosed = errors.New("redis: client is closed")
var ErrPoolTimeout = errors.New("redis: connection pool timeout")

var timers = sync.Pool{
	New: func() interface{} {
		t := time.NewTimer(time.Hour)
		t.Stop()
		return t
	},
}

// Stats contains pool state information and accumulated stats.
type Stats struct {
	Hits     uint32 // number of times free connection was found in the pool
	Misses   uint32 // number of times free connection was NOT found in the pool
	Timeouts uint32 // number of times a wait timeout occurred

	TotalConns uint32 // number of total connections in the pool
	IdleConns  uint32 // number of idle connections in the pool
	StaleConns uint32 // number of stale connections removed from the pool
}

type Pooler interface {
	NewConn(context.Context) (*Conn, error)
	CloseConn(*Conn) error

	Get(context.Context) (*Conn, error)
	Put(*Conn)
	Remove(*Conn, error)

	Len() int
	IdleLen() int
	Stats() *Stats

	Close() error
}

type Options struct {
	Dialer  func(context.Context) (net.Conn, error)
	OnClose func(*Conn) error

	PoolSize           int
	MinIdleConns       int
	MaxConnAge         time.Duration
	PoolTimeout        time.Duration
	IdleTimeout        time.Duration
	IdleCheckFrequency time.Duration
}

type lastDialErrorWrap struct {
	err error
}

type ConnPool struct {
	opt *Options

	dialErrorsNum uint32 // atomic

	lastDialError atomic.Value

	queue chan struct{}

	connsMu      sync.Mutex
	conns        []*Conn
	idleConns    []*Conn
	poolSize     int
	idleConnsLen int

	stats Stats

	_closed  uint32 // atomic
	closedCh chan struct{}
}

var _ Pooler = (*ConnPool)(nil)

func NewConnPool(opt *Options) *ConnPool {
	p := &ConnPool{
		opt: opt,

		queue:     make(chan struct{}, opt.PoolSize),
		conns:     make([]*Conn, 0, opt.PoolSize),
		idleConns: make([]*Conn, 0, opt.PoolSize),
		closedCh:  make(chan struct{}),
	}

	p.connsMu.Lock()
	p.checkMinIdleConns()
	p.connsMu.Unlock()

	if opt.IdleTimeout > 0 && opt.IdleCheckFrequency > 0 {
		go p.reaper(opt.IdleCheckFrequency)
	}

	return p
}

func (p *ConnPool) checkMinIdleConns() {
	if p.opt.MinIdleConns == 0 {
		return
	}
	for p.poolSize < p.opt.PoolSize && p.idleConnsLen < p.opt.MinIdleConns {
		p.poolSize++
		p.idleConnsLen++
		go func() {
			err := p.addIdleConn()
			if err != nil {
				p.connsMu.Lock()
				p.poolSize--
				p.idleConnsLen--
				p.connsMu.Unlock()
			}
		}()
	}
}

func (p *ConnPool) addIdleConn() error {
	cn, err := p.dialConn(context.TODO(), true)
	if err != nil {
		return err
	}

	p.connsMu.Lock()
	p.conns = append(p.conns, cn)
	p.idleConns = append(p.idleConns, cn)
	p.connsMu.Unlock()
	return nil
}

func (p *ConnPool) NewConn(ctx context.Context) (*Conn, error) {
	return p.newConn(ctx, false)
}

func (p *ConnPool) newConn(ctx context.Context, pooled bool) (*Conn, error) {
	cn, err := p.dialConn(ctx, pooled)
	if err != nil {
		return nil, err
	}

	p.connsMu.Lock()
	p.conns = append(p.conns, cn)
	if pooled {
		// If pool is full remove the cn on next Put.
		if p.poolSize >= p.opt.PoolSize {
			cn.pooled = false
		} else {
			p.poolSize++
		}
	}
	p.connsMu.Unlock()
	return cn, nil
}

func (p *ConnPool) dialConn(ctx context.Context, pooled bool) (*Conn, error) {
	if p.closed() {
		return nil, ErrClosed
	}

	if atomic.LoadUint32(&p.dialErrorsNum) >= uint32(p.opt.PoolSize) {
		return nil, p.getLastDialError()
	}

	netConn, err := p.opt.Dialer(ctx)
	if err != nil {
		p.setLastDialError(err)
		if atomic.AddUint32(&p.dialErrorsNum, 1) == uint32(p.opt.PoolSize) {
			go p.tryDial()
		}
		return nil, err
	}

	cn := NewConn(netConn)
	cn.pooled = pooled
	return cn, nil
}

func (p *ConnPool) tryDial() {
	for {
		if p.closed() {
			return
		}

		conn, err := p.opt.Dialer(context.Background())
		if err != nil {
			p.setLastDialError(err)
			time.Sleep(time.Second)
			continue
		}

		atomic.StoreUint32(&p.dialErrorsNum, 0)
		_ = conn.Close()
		return
	}
}

func (p *ConnPool) setLastDialError(err error) {
	p.lastDialError.Store(&lastDialErrorWrap{err: err})
}

func (p *ConnPool) getLastDialError() error {
	err, _ := p.lastDialError.Load().(*lastDialErrorWrap)
	if err != nil {
		return err.err
	}
	return nil
}

// Get returns existed connection from the pool or creates a new one.
func (p *ConnPool) Get(ctx context.Context) (*Conn, error) {
	if p.closed() {
		return nil, ErrClosed
	}

	err := p.waitTurn(ctx)
	if err != nil {
		return nil, err
	}

	for {
		p.connsMu.Lock()
		cn := p.popIdle()
		p.connsMu.Unlock()

		if cn == nil {
			break
		}

		if p.isStaleConn(cn) {
			_ = p.CloseConn(cn)
			continue
		}

		atomic.AddUint32(&p.stats.Hits, 1)
		return cn, nil
	}

	atomic.AddUint32(&p.stats.Misses, 1)

	newcn, err := p.newConn(ctx, true)
	if err != nil {
		p.freeTurn()
		return nil, err
	}

	return newcn, nil
}

func (p *ConnPool) getTurn() {
	p.queue <- struct{}{}
}

func (p *ConnPool) waitTurn(ctx context.Context) error {
	select {
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	select {
	case p.queue <- struct{}{}:
		return nil
	default:
	}

	timer := timers.Get().(*time.Timer)
	timer.Reset(p.opt.PoolTimeout)

	select {
	case <-ctx.Done():
		if !timer.Stop() {
			<-timer.C
		}
		timers.Put(timer)
		return ctx.Err()
	case p.queue <- struct{}{}:
		if !timer.Stop() {
			<-timer.C
		}
		timers.Put(timer)
		return nil
	case <-timer.C:
		timers.Put(timer)
		atomic.AddUint32(&p.stats.Timeouts, 1)
		return ErrPoolTimeout
	}
}

func (p *ConnPool) freeTurn() {
	<-p.queue
}

func (p *ConnPool) popIdle() *Conn {
	if len(p.idleConns) == 0 {
		return nil
	}

	idx := len(p.idleConns) - 1
	cn := p.idleConns[idx]
	p.idleConns = p.idleConns[:idx]
	p.idleConnsLen--
	p.checkMinIdleConns()
	return cn
}

func (p *ConnPool) Put(cn *Conn) {
	if cn.rd.Buffered() > 0 {
		internal.Logger.Printf("Conn has unread data")
		p.Remove(cn, BadConnError{})
		return
	}

	if !cn.pooled {
		p.Remove(cn, nil)
		return
	}

	p.connsMu.Lock()
	p.idleConns = append(p.idleConns, cn)
	p.idleConnsLen++
	p.connsMu.Unlock()
	p.freeTurn()
}

func (p *ConnPool) Remove(cn *Conn, reason error) {
	p.removeConnWithLock(cn)
	p.freeTurn()
	_ = p.closeConn(cn)
}

func (p *ConnPool) CloseConn(cn *Conn) error {
	p.removeConnWithLock(cn)
	return p.closeConn(cn)
}

func (p *ConnPool) removeConnWithLock(cn *Conn) {
	p.connsMu.Lock()
	p.removeConn(cn)
	p.connsMu.Unlock()
}

func (p *ConnPool) removeConn(cn *Conn) {
	for i, c := range p.conns {
		if c == cn {
			p.conns = append(p.conns[:i], p.conns[i+1:]...)
			if cn.pooled {
				p.poolSize--
				p.checkMinIdleConns()
			}
			return
		}
	}
}

func (p *ConnPool) closeConn(cn *Conn) error {
	if p.opt.OnClose != nil {
		_ = p.opt.OnClose(cn)
	}
	return cn.Close()
}

// Len returns total number of connections.
func (p *ConnPool) Len() int {
	p.connsMu.Lock()
	n := len(p.conns)
	p.connsMu.Unlock()
	return n
}

// IdleLen returns number of idle connections.
func (p *ConnPool) IdleLen() int {
	p.connsMu.Lock()
	n := p.idleConnsLen
	p.connsMu.Unlock()
	return n
}

func (p *ConnPool) Stats() *Stats {
	idleLen := p.IdleLen()
	return &Stats{
		Hits:     atomic.LoadUint32(&p.stats.Hits),
		Misses:   atomic.LoadUint32(&p.stats.Misses),
		Timeouts: atomic.LoadUint32(&p.stats.Timeouts),

		TotalConns: uint32(p.Len()),
		IdleConns:  uint32(idleLen),
		StaleConns: atomic.LoadUint32(&p.stats.StaleConns),
	}
}

func (p *ConnPool) closed() bool {
	return atomic.LoadUint32(&p._closed) == 1
}

func (p *ConnPool) Filter(fn func(*Conn) bool) error {
	var firstErr error
	p.connsMu.Lock()
	for _, cn := range p.conns {
		if fn(cn) {
			if err := p.closeConn(cn); err != nil && firstErr == nil {
				firstErr = err
			}
		}
	}
	p.connsMu.Unlock()
	return firstErr
}

func (p *ConnPool) Close() error {
	if !atomic.CompareAndSwapUint32(&p._closed, 0, 1) {
		return ErrClosed
	}
	close(p.closedCh)

	var firstErr error
	p.connsMu.Lock()
	for _, cn := range p.conns {
		if err := p.closeConn(cn); err != nil && firstErr == nil {
			firstErr = err
		}
	}
	p.conns = nil
	p.poolSize = 0
	p.idleConns = nil
	p.idleConnsLen = 0
	p.connsMu.Unlock()

	return firstErr
}

func (p *ConnPool) reaper(frequency time.Duration) {
	ticker := time.NewTicker(frequency)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			// It is possible that ticker and closedCh arrive together,
			// and select pseudo-randomly pick ticker case, we double
			// check here to prevent being executed after closed.
			if p.closed() {
				return
			}
			_, err := p.ReapStaleConns()
			if err != nil {
				internal.Logger.Printf("ReapStaleConns failed: %s", err)
				continue
			}
		case <-p.closedCh:
			return
		}
	}
}

func (p *ConnPool) ReapStaleConns() (int, error) {
	var n int
	for {
		p.getTurn()

		p.connsMu.Lock()
		cn := p.reapStaleConn()
		p.connsMu.Unlock()
		p.freeTurn()

		if cn != nil {
			_ = p.closeConn(cn)
			n++
		} else {
			break
		}
	}
	atomic.AddUint32(&p.stats.StaleConns, uint32(n))
	return n, nil
}

func (p *ConnPool) reapStaleConn() *Conn {
	if len(p.idleConns) == 0 {
		return nil
	}

	cn := p.idleConns[0]
	if !p.isStaleConn(cn) {
		return nil
	}

	p.idleConns = append(p.idleConns[:0], p.idleConns[1:]...)
	p.idleConnsLen--
	p.removeConn(cn)

	return cn
}

func (p *ConnPool) isStaleConn(cn *Conn) bool {
	if p.opt.IdleTimeout == 0 && p.opt.MaxConnAge == 0 {
		return false
	}

	now := time.Now()
	if p.opt.IdleTimeout > 0 && now.Sub(cn.UsedAt()) >= p.opt.IdleTimeout {
		return true
	}
	if p.opt.MaxConnAge > 0 && now.Sub(cn.createdAt) >= p.opt.MaxConnAge {
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package bencode

import (
	"bytes"
	"reflect"
	"testing"
)

type testcase struct {
	in  string
	out interface{}
}

func TestConsumeValue(t *testing.T) {
	tests := []testcase{
		testcase{"i2e", 2},
		testcase{"1:a", "a"},
		testcase{"li2ei42ei666ee", []int{2, 42, 666}},
		testcase{"d1:Ai42e1:B3:xyze", struct {
			A int
			B string
		}{A: 42, B: "xyz"}},
	}
	for _, test := range tests {
		out := reflect.New(reflect.TypeOf(test.out)).Elem()
		consumeValue(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}

func TestConsumeInt(t *testing.T) {
	tests := []testcase{
		testcase{"i2e", 2},
		testcase{"i42e", 42},
		testcase{"i666e", 666},
	}
	for _, test := range tests {
		var o int
		out := reflect.ValueOf(&o).Elem()
		consumeInt(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}

func TestConsumeString(t *testing.T) {
	tests := []testcase{
		testcase{"1:s", "s"},
		testcase{"4:butt", "butt"},
		testcase{"11:buttfartass", "buttfartass"},
	}
	for _, test := range tests {
		var o string
		out := reflect.ValueOf(&o).Elem()
		consumeString(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}

func TestConsumeList(t *testing.T) {
	tests := []testcase{
		testcase{"li2ei42ei666ee", []int{2, 42, 666}},
		testcase{"l1:a1:b1:ce", []string{"a", "b", "c"}},
		testcase{"lli2eeli42eee", [][]int{[]int{2}, []int{42}}},
	}

	for _, test := range tests {
		out := reflect.New(reflect.TypeOf(test.out)).Elem()
		consumeList(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}

func TestConsumeDict(t *testing.T) {
	tests := []testcase{
		testcase{"d1:Ai42e1:B3:xyze", struct {
			A int
			B string
		}{A: 42, B: "xyz"}},
	}

	for _, test := range tests {
		out := reflect.New(reflect.TypeOf(test.out)).Elem()
		consumeDict(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}
<commit_msg>Adding list of dicts testcase<commit_after>package bencode

import (
	"bytes"
	"reflect"
	"testing"
)

type testcase struct {
	in  string
	out interface{}
}

func TestConsumeValue(t *testing.T) {
	tests := []testcase{
		testcase{"i2e", 2},
		testcase{"1:a", "a"},
		testcase{"li2ei42ei666ee", []int{2, 42, 666}},
		testcase{"d1:Ai42e1:B3:xyze", struct {
			A int
			B string
		}{A: 42, B: "xyz"}},
	}
	for _, test := range tests {
		out := reflect.New(reflect.TypeOf(test.out)).Elem()
		consumeValue(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}

func TestConsumeInt(t *testing.T) {
	tests := []testcase{
		testcase{"i2e", 2},
		testcase{"i42e", 42},
		testcase{"i666e", 666},
	}
	for _, test := range tests {
		var o int
		out := reflect.ValueOf(&o).Elem()
		consumeInt(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}

func TestConsumeString(t *testing.T) {
	tests := []testcase{
		testcase{"1:s", "s"},
		testcase{"4:butt", "butt"},
		testcase{"11:buttfartass", "buttfartass"},
	}
	for _, test := range tests {
		var o string
		out := reflect.ValueOf(&o).Elem()
		consumeString(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}

func TestConsumeList(t *testing.T) {
	tests := []testcase{
		testcase{"li2ei42ei666ee", []int{2, 42, 666}},
		testcase{"l1:a1:b1:ce", []string{"a", "b", "c"}},
		testcase{"lli2eeli42eee", [][]int{[]int{2}, []int{42}}},
		testcase{"ld1:A4:butted1:A4:fartee", []struct{ A string }{
			struct{ A string }{"butt"},
			struct{ A string }{"fart"},
		}},
	}

	for _, test := range tests {
		out := reflect.New(reflect.TypeOf(test.out)).Elem()
		consumeList(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}

func TestConsumeDict(t *testing.T) {
	tests := []testcase{
		testcase{"d1:Ai42e1:B3:xyze", struct {
			A int
			B string
		}{A: 42, B: "xyz"}},
	}

	for _, test := range tests {
		out := reflect.New(reflect.TypeOf(test.out)).Elem()
		consumeDict(out, bytes.NewBuffer([]byte(test.in)))
		if i := out.Interface(); !reflect.DeepEqual(i, test.out) {
			t.Error("Expecting", test.out, "got", i)
		}
	}
}
<|endoftext|>


<commit_before>// web100lib provides Go bindings to some functions in the web100 library.
package main

// Cgo directives must immediately preceed 'import "C"' below.

/*
#include <stdio.h>
#include <stdlib.h>
#include <web100.h>
#include <web100-int.h>

#include <arpa/inet.h>

void print_bytes(size_t var_size, void *var_data) {
	unsigned char *data = (unsigned char *)var_data;
	int i = 0;
	fflush(stdout);
	for (i = 0; i < var_size; i++ ) {
		fprintf(stdout, "%02x ", data[i]);
	}
	fflush(stdout);
}
*/
import "C"

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"unsafe"
	//"github.com/kr/pretty"
)

var (
	filename = flag.String("filename", "", "Trace filename.")
)

// Necessary web100 functions:
//  + web100_log_open_read(filename)
//  + web100_log_close_read(log_)
//  + snap_ = web100_snapshot_alloc_from_log(log_);
//  + web100_snap_from_log(snap_, log_)
//
//  + for (web100_var *var = web100_var_head(group_);
//  +      var != NULL;
//  +      var = web100_var_next(var)) {
//
//   web100_get_log_agent(log_)
//   web100_get_log_time(log_);
//   + web100_get_log_group(log_);
//
//   connection_ = web100_get_log_connection(log_);

// Notes:
//  - See: https://golang.org/cmd/cgo/#hdr-Go_references_to_C
//
// Discoveries:
//  - Not all C macros exist in the "C" namespace.
//  - 'NULL' is usually equivalent to 'nil'

// Web100 maintains state associated with a web100 log file.
type Web100 struct {
	// Do not export unsafe pointers.
	log  unsafe.Pointer
	snap unsafe.Pointer
}

// Open prepares a web100 log file for reading. The caller must call Close on
// the returned Web100 instance to release resources.
func Open(filename string) (*Web100, error) {
	c_filename := C.CString(filename)
	defer C.free(unsafe.Pointer(c_filename))

	log := C.web100_log_open_read(c_filename)
	if log == nil {
		return nil, fmt.Errorf(C.GoString(C.web100_strerror(C.web100_errno)))
	}

	// Pre-allocate a snapshot record.
	snap := C.web100_snapshot_alloc_from_log(log)

	w := &Web100{
		log:  unsafe.Pointer(log),
		snap: unsafe.Pointer(snap),
	}
	return w, nil
}

// Next iterates through the web100 log file and returns the next snapshot
// record in the form of a map.
func (w *Web100) Next() error {
	log := (*C.web100_log)(w.log)
	snap := (*C.web100_snapshot)(w.snap)

	// Read the next web100_snaplog data from underlying file.
	err := C.web100_snap_from_log(snap, log)
	if err == C.EOF {
		return io.EOF
	}
	if err != C.WEB100_ERR_SUCCESS {
		return fmt.Errorf(C.GoString(C.web100_strerror(err)))
	}
	return nil
}

// LogValues returns a map of values from the web100 log. IPv6 address
// connection information is not available.
func (w *Web100) LogValues() (map[string]string, error) {
	log := (*C.web100_log)(w.log)

	agent := C.web100_get_log_agent(log)

	results := make(map[string]string)
	results["web100_log_entry.version"] = C.GoString(C.web100_get_agent_version(agent))

	time := C.web100_get_log_time(log)
	results["web100_log_entry.log_time"] = fmt.Sprintf("%d", int64(time))

	conn := C.web100_get_log_connection(log)
	// NOTE: web100_connection_spec_v6 is not filled in by the web100 library.
	// NOTE: addrtype is always WEB100_ADDRTYPE_UNKNOWN.
	results["web100_log_entry.connection_spec.local_af"] = ""
	var spec C.struct_web100_connection_spec
	C.web100_get_connection_spec(conn, &spec)

	addr := C.struct_in_addr{C.in_addr_t(spec.src_addr)}
	results["web100_log_entry.connection_spec.local_ip"] = C.GoString(C.inet_ntoa(addr))
	results["web100_log_entry.connection_spec.local_port"] = fmt.Sprintf("%d", spec.src_port)

	addr = C.struct_in_addr{C.in_addr_t(spec.dst_addr)}
	results["web100_log_entry.connection_spec.remote_ip"] = C.GoString(C.inet_ntoa(addr))
	results["web100_log_entry.connection_spec.remote_port"] = fmt.Sprintf("%d", spec.dst_port)

	return results, nil
}

// SnapValues converts all variables in the latest snap record into a results
// map.
func (w *Web100) SnapValues() (map[string]string, error) {
	log := (*C.web100_log)(w.log)
	snap := (*C.web100_snapshot)(w.snap)

	results := make(map[string]string)

	var_text := C.calloc(2*C.WEB100_VALUE_LEN_MAX, 1) // Use a better size.
	defer C.free(var_text)

	var_data := C.calloc(C.WEB100_VALUE_LEN_MAX, 1)
	defer C.free(var_data)

	// Parses variables from most recent web100_snapshot data.
	group := C.web100_get_log_group(log)
	for v := C.web100_var_head(group); v != nil; v = C.web100_var_next(v) {

		name := C.web100_get_var_name(v)
		var_size := C.web100_get_var_size(v)
		var_type := C.web100_get_var_type(v)

		// Read the raw variable data from the snapshot data.
		err := C.web100_snap_read(v, snap, var_data)
		if err != C.WEB100_ERR_SUCCESS {
			return nil, fmt.Errorf(C.GoString(C.web100_strerror(err)))
		}

		// Convert raw var_data into a string based on var_type.
		C.web100_value_to_textn((*C.char)(var_text), C.WEB100_VALUE_LEN_MAX, (C.WEB100_TYPE)(var_type), var_data)
		results[fmt.Sprintf("web100_log_entry.snap.%s", C.GoString(name))] = C.GoString((*C.char)(var_text))

		fmt.Printf("name: %-20s type: %d %d size %d: %-30s ", C.GoString(name), C.WEB100_TYPE_INTEGER32, var_type, var_size,
			C.GoString((*C.char)(var_text)))
		C.print_bytes(var_size, var_data)
		fmt.Printf("\n")
	}

	return results, nil
}

// Close releases resources created by Open.
func (w *Web100) Close() error {
	snap := (*C.web100_snapshot)(w.snap)
	C.web100_snapshot_free(snap)

	log := (*C.web100_log)(w.log)
	err := C.web100_log_close_read(log)
	if err != C.WEB100_ERR_SUCCESS {
		return fmt.Errorf(C.GoString(C.web100_strerror(err)))
	}

	// Clear pointer after free.
	w.log = nil
	w.snap = nil
	return nil
}

func LookupError(errnum int) string {
	return C.GoString(C.web100_strerror(C.int(errnum)))
}

func PrettyPrint(results map[string]string) {
	b, err := json.MarshalIndent(results, "", "  ")
	if err != nil {
		fmt.Println("error:", err)
	}
	fmt.Print(string(b))
}

func main() {
	flag.Parse()

	fmt.Println(LookupError(0))
	w, err := Open(*filename)
	if err != nil {
		panic(err)
	}
	fmt.Printf("%#v\n", w)

	results, err := w.LogValues()
	if err != nil {
		panic(err)
	}
	PrettyPrint(results)

	// Find and print the last web100 snapshot record.
	for {
		err = w.Next()
		if err != nil {
			break
		}
	}
	if err != io.EOF {
		panic(err)
	}
	results, err = w.SnapValues()
	if err != nil {
		panic(err)
	}
	PrettyPrint(results)
	w.Close()
	fmt.Printf("%#v\n", w)
}
<commit_msg>Rename.<commit_after><|endoftext|>