<commit_before>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains high level helper functions and easy entry points for the
// entire discordgo package.  These functions are being developed and are very
// experimental at this point.  They will most likely change so please use the
// low level functions if that's a problem.

// Package discordgo provides Discord binding for Go
package discordgo

import (
	"errors"
	"fmt"
	"net/http"
	"runtime"
	"time"
)

// VERSION of DiscordGo, follows Semantic Versioning. (http://semver.org/)
const VERSION = "0.22.0"

// ErrMFA will be risen by New when the user has 2FA.
var ErrMFA = errors.New("account has 2FA enabled")

// New creates a new Discord session and will automate some startup
// tasks if given enough information to do so.  Currently you can pass zero
// arguments and it will return an empty Discord session.
// There are 3 ways to call New:
//     With a single auth token - All requests will use the token blindly
//         (just tossing it into the HTTP Authorization header);
//         no verification of the token will be done and requests may fail.
//         IF THE TOKEN IS FOR A BOT, IT MUST BE PREFIXED WITH `BOT `
//         eg: `"Bot <token>"`
//         IF IT IS AN OAUTH2 ACCESS TOKEN, IT MUST BE PREFIXED WITH `Bearer `
//         eg: `"Bearer <token>"`
//     With an email and password - Discord will sign in with the provided
//         credentials.
//     With an email, password and auth token - Discord will verify the auth
//         token, if it is invalid it will sign in with the provided
//         credentials. This is the Discord recommended way to sign in.
//
// NOTE: While email/pass authentication is supported by DiscordGo it is
// HIGHLY DISCOURAGED by Discord. Please only use email/pass to obtain a token
// and then use that authentication token for all future connections.
// Also, doing any form of automation with a user (non Bot) account may result
// in that account being permanently banned from Discord.
func New(args ...interface{}) (s *Session, err error) {

	// Create an empty Session interface.
	s = &Session{
		State:                  NewState(),
		Ratelimiter:            NewRatelimiter(),
		StateEnabled:           true,
		Compress:               true,
		ShouldReconnectOnError: true,
		ShardID:                0,
		ShardCount:             1,
		MaxRestRetries:         3,
		Client:                 &http.Client{Timeout: (20 * time.Second)},
		UserAgent:              "DiscordBot (https://github.com/bwmarrin/discordgo, v" + VERSION + ")",
		sequence:               new(int64),
		LastHeartbeatAck:       time.Now().UTC(),
	}

	// Initilize the Identify Package with defaults
	// These can be modified prior to calling Open()
	s.Identify.Compress = true
	s.Identify.LargeThreshold = 250
	s.Identify.GuildSubscriptions = true
	s.Identify.Properties.OS = runtime.GOOS
	s.Identify.Properties.Browser = "DiscordGo v" + VERSION
	s.Identify.Intents = MakeIntent(IntentsAllWithoutPrivileged)

	// If no arguments are passed return the empty Session interface.
	if args == nil {
		return
	}

	// Variables used below when parsing func arguments
	var auth, pass string

	// Parse passed arguments
	for _, arg := range args {

		switch v := arg.(type) {

		case []string:
			if len(v) > 3 {
				err = fmt.Errorf("too many string parameters provided")
				return
			}

			// First string is either token or username
			if len(v) > 0 {
				auth = v[0]
			}

			// If second string exists, it must be a password.
			if len(v) > 1 {
				pass = v[1]
			}

			// If third string exists, it must be an auth token.
			if len(v) > 2 {
				s.Identify.Token = v[2]
				s.Token = v[2] // TODO: Remove, Deprecated - Kept for backwards compatibility.
			}

		case string:
			// First string must be either auth token or username.
			// Second string must be a password.
			// Only 2 input strings are supported.

			if auth == "" {
				auth = v
			} else if pass == "" {
				pass = v
			} else if s.Token == "" {
				s.Identify.Token = v
				s.Token = v // TODO: Remove, Deprecated - Kept for backwards compatibility.
			} else {
				err = fmt.Errorf("too many string parameters provided")
				return
			}

			//		case Config:
			// TODO: Parse configuration struct

		default:
			err = fmt.Errorf("unsupported parameter type provided")
			return
		}
	}

	// If only one string was provided, assume it is an auth token.
	// Otherwise get auth token from Discord, if a token was specified
	// Discord will verify it for free, or log the user in if it is
	// invalid.
	if pass == "" {
		s.Identify.Token = auth
		s.Token = auth // TODO: Remove, Deprecated - Kept for backwards compatibility.
	} else {
		err = s.Login(auth, pass)
		// TODO: Remove last s.Token part, Deprecated - Kept for backwards compatibility.
		if err != nil || s.Identify.Token == "" || s.Token == "" {
			if s.MFA {
				err = ErrMFA
			} else {
				err = fmt.Errorf("Unable to fetch discord authentication token. %v", err)
			}
			return
		}
	}

	return
}
<commit_msg>Change library version to alpha for v8<commit_after>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains high level helper functions and easy entry points for the
// entire discordgo package.  These functions are being developed and are very
// experimental at this point.  They will most likely change so please use the
// low level functions if that's a problem.

// Package discordgo provides Discord binding for Go
package discordgo

import (
	"errors"
	"fmt"
	"net/http"
	"runtime"
	"time"
)

// VERSION of DiscordGo, follows Semantic Versioning. (http://semver.org/)
const VERSION = "0.23.0-v8alpha"

// ErrMFA will be risen by New when the user has 2FA.
var ErrMFA = errors.New("account has 2FA enabled")

// New creates a new Discord session and will automate some startup
// tasks if given enough information to do so.  Currently you can pass zero
// arguments and it will return an empty Discord session.
// There are 3 ways to call New:
//     With a single auth token - All requests will use the token blindly
//         (just tossing it into the HTTP Authorization header);
//         no verification of the token will be done and requests may fail.
//         IF THE TOKEN IS FOR A BOT, IT MUST BE PREFIXED WITH `BOT `
//         eg: `"Bot <token>"`
//         IF IT IS AN OAUTH2 ACCESS TOKEN, IT MUST BE PREFIXED WITH `Bearer `
//         eg: `"Bearer <token>"`
//     With an email and password - Discord will sign in with the provided
//         credentials.
//     With an email, password and auth token - Discord will verify the auth
//         token, if it is invalid it will sign in with the provided
//         credentials. This is the Discord recommended way to sign in.
//
// NOTE: While email/pass authentication is supported by DiscordGo it is
// HIGHLY DISCOURAGED by Discord. Please only use email/pass to obtain a token
// and then use that authentication token for all future connections.
// Also, doing any form of automation with a user (non Bot) account may result
// in that account being permanently banned from Discord.
func New(args ...interface{}) (s *Session, err error) {

	// Create an empty Session interface.
	s = &Session{
		State:                  NewState(),
		Ratelimiter:            NewRatelimiter(),
		StateEnabled:           true,
		Compress:               true,
		ShouldReconnectOnError: true,
		ShardID:                0,
		ShardCount:             1,
		MaxRestRetries:         3,
		Client:                 &http.Client{Timeout: (20 * time.Second)},
		UserAgent:              "DiscordBot (https://github.com/bwmarrin/discordgo, v" + VERSION + ")",
		sequence:               new(int64),
		LastHeartbeatAck:       time.Now().UTC(),
	}

	// Initilize the Identify Package with defaults
	// These can be modified prior to calling Open()
	s.Identify.Compress = true
	s.Identify.LargeThreshold = 250
	s.Identify.GuildSubscriptions = true
	s.Identify.Properties.OS = runtime.GOOS
	s.Identify.Properties.Browser = "DiscordGo v" + VERSION
	s.Identify.Intents = MakeIntent(IntentsAllWithoutPrivileged)

	// If no arguments are passed return the empty Session interface.
	if args == nil {
		return
	}

	// Variables used below when parsing func arguments
	var auth, pass string

	// Parse passed arguments
	for _, arg := range args {

		switch v := arg.(type) {

		case []string:
			if len(v) > 3 {
				err = fmt.Errorf("too many string parameters provided")
				return
			}

			// First string is either token or username
			if len(v) > 0 {
				auth = v[0]
			}

			// If second string exists, it must be a password.
			if len(v) > 1 {
				pass = v[1]
			}

			// If third string exists, it must be an auth token.
			if len(v) > 2 {
				s.Identify.Token = v[2]
				s.Token = v[2] // TODO: Remove, Deprecated - Kept for backwards compatibility.
			}

		case string:
			// First string must be either auth token or username.
			// Second string must be a password.
			// Only 2 input strings are supported.

			if auth == "" {
				auth = v
			} else if pass == "" {
				pass = v
			} else if s.Token == "" {
				s.Identify.Token = v
				s.Token = v // TODO: Remove, Deprecated - Kept for backwards compatibility.
			} else {
				err = fmt.Errorf("too many string parameters provided")
				return
			}

			//		case Config:
			// TODO: Parse configuration struct

		default:
			err = fmt.Errorf("unsupported parameter type provided")
			return
		}
	}

	// If only one string was provided, assume it is an auth token.
	// Otherwise get auth token from Discord, if a token was specified
	// Discord will verify it for free, or log the user in if it is
	// invalid.
	if pass == "" {
		s.Identify.Token = auth
		s.Token = auth // TODO: Remove, Deprecated - Kept for backwards compatibility.
	} else {
		err = s.Login(auth, pass)
		// TODO: Remove last s.Token part, Deprecated - Kept for backwards compatibility.
		if err != nil || s.Identify.Token == "" || s.Token == "" {
			if s.MFA {
				err = ErrMFA
			} else {
				err = fmt.Errorf("Unable to fetch discord authentication token. %v", err)
			}
			return
		}
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"

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

func Extract(url string) ([]string, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return nil, fmt.Errorf("getting %s: %s", url, resp.Status)
	}

	doc, err := html.Parse(resp.Body)
	resp.Body.Close()
	if err != nil {
		return nil, fmt.Errorf("parsing %s as HTML: %v", url, err)
	}

	var links []string
	visitNode := func(n *html.Node) {
		if n.Type == html.ElementNode && n.Data == "a" {
			for _, a := range n.Attr {
				if a.Key != "href" {
					continue
				}
				link, err := resp.Request.URL.Parse(a.Val)
				if err != nil {
					continue // ignore bad URLs
				}
				links = append(links, link.String())
			}
		}
	}
	forEachNode(doc, visitNode, nil)
	return links, nil
}

func forEachNode(n *html.Node, pre, post func(n *html.Node)) {
	if pre != nil {
		pre(n)
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		forEachNode(c, pre, post)
	}
	if post != nil {
		post(n)
	}
}
<commit_msg>[5.13] Workaround for js href.<commit_after>package main

import (
	"fmt"
	"net/http"
	"strings"

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

func Extract(url string) ([]string, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		resp.Body.Close()
		return nil, fmt.Errorf("getting %s: %s", url, resp.Status)
	}

	doc, err := html.Parse(resp.Body)
	resp.Body.Close()
	if err != nil {
		return nil, fmt.Errorf("parsing %s as HTML: %v", url, err)
	}

	var links []string
	visitNode := func(n *html.Node) {
		if n.Type == html.ElementNode && n.Data == "a" {
			for _, a := range n.Attr {
				if a.Key != "href" {
					continue
				}
				if strings.HasPrefix(a.Val, "javascript:") {
					continue
				}
				link, err := resp.Request.URL.Parse(a.Val)
				if err != nil {
					continue // ignore bad URLs
				}
				links = append(links, link.String())
			}
		}
	}
	forEachNode(doc, visitNode, nil)
	return links, nil
}

func forEachNode(n *html.Node, pre, post func(n *html.Node)) {
	if pre != nil {
		pre(n)
	}
	for c := n.FirstChild; c != nil; c = c.NextSibling {
		forEachNode(c, pre, post)
	}
	if post != nil {
		post(n)
	}
}
<|endoftext|>


<commit_before>package game

import (
	"sync"
	"time"
	"runtime"

	"buildblast/lib/coords"
	"buildblast/lib/mapgen"
)

type ChunkGenerator struct {
	// Chunks are sent to this channel as they are generated
	Generated chan ChunkGenerationResult

	chunks    map[coords.Chunk]ChunkStatus
	mutex     sync.Mutex
	goPool    chan bool
	generator mapgen.Generator
}

type ChunkGenerationResult struct {
	cc    coords.Chunk
	chunk *mapgen.Chunk
}

type ChunkStatus struct {
	generating bool
	generated  bool
	priority   int
}

func NewChunkGenerator(generator mapgen.Generator) *ChunkGenerator {
	cm := new(ChunkGenerator)
	cm.chunks = make(map[coords.Chunk]ChunkStatus, 10)
	cm.Generated = make(chan ChunkGenerationResult, 10)
	// Only cores - 1 chunks can be generated at a time.
	maxActiveThreads := runtime.NumCPU() - 1
	if maxActiveThreads < 1 {
		maxActiveThreads = 1
	}
	cm.goPool = make(chan bool, maxActiveThreads)
	for i := 0; i < maxActiveThreads; i++ {
		cm.goPool <- true
	}
	cm.generator = generator
	return cm
}

// cm.Lock() MUST BE HELD by the caller, or else calling
// this function is unsafe.
func (cm *ChunkGenerator) queue(cc coords.Chunk, priority int) {
	status := cm.chunks[cc]
	if status.generated || status.generating {
		return
	}
	status.priority += priority
	cm.chunks[cc] = status
}

func (cm *ChunkGenerator) QueueChunksNearby(wc coords.World) {
	cm.mutex.Lock()
	defer cm.mutex.Unlock()

	queue := func(cc coords.Chunk, priority int) {
		cm.queue(cc, priority)
	}

	EachChunkNearby(wc, queue)
}

func (cm *ChunkGenerator) Top() (cc coords.Chunk, valid bool) {
	cm.mutex.Lock()
	defer cm.mutex.Unlock()

	highest := -1
	for key, val := range cm.chunks {
		if val.priority > highest && !val.generated && !val.generating {
			highest = val.priority
			cc = key
		}
	}
	if highest != -1 {
		return cc, true
	}
	return cc, false
}

func (cm *ChunkGenerator) Run() {
	for {
		cc, valid := cm.Top()
		if !valid {
			<-time.After(time.Second / 60)
			continue
		}

		cm.mutex.Lock()
		status := cm.chunks[cc]
		status.generating = true
		cm.chunks[cc] = status
		cm.mutex.Unlock()

		go cm.generate(cc)
	}
}

func (cm *ChunkGenerator) generate(cc coords.Chunk) {
	// Attempt to use a slot
	<- cm.goPool
	chunk := cm.generator.Chunk(cc)

	cm.Generated <- ChunkGenerationResult{
		cc:    cc,
		chunk: chunk,
	}

	cm.mutex.Lock()
	status := cm.chunks[cc]
	status.generating = false
	status.generated = true
	cm.chunks[cc] = status
	cm.mutex.Unlock()
	// Free our slot
	cm.goPool <- true
}

func EachChunkNearby(wc coords.World, cb func(cc coords.Chunk, priority int)) {
	occ := func(cc coords.Chunk, x, y, z int) coords.Chunk {
		return coords.Chunk{
			X: cc.X + x,
			Y: cc.Y + y,
			Z: cc.Z + z,
		}
	}

	eachWithin := func(cc coords.Chunk, xdist, ydist, zdist int, cb func(newCC coords.Chunk, dist int)) {
		abs := func(n int) int {
			if n < 0 {
				return -n
			}
			return n
		}
		dist := func(x, y, z int) int {
			return abs(x) + abs(y) + abs(z)
		}

		cb(cc, 0)
		for x := -xdist; x <= xdist; x++ {
			for y := -ydist; y <= ydist; y++ {
				for z := -zdist; z <= zdist; z++ {
					cb(occ(cc, x, y, z), dist(x, y, z))
				}
			}
		}
	}

	cc := wc.Chunk()
	eachWithin(cc, 2, 1, 2, func(newCC coords.Chunk, dist int) {
		// We want to prioritize further away chunks lower, but the
		// priority must be a positive integer.
		cb(newCC, 10-dist)
	})

	oc := wc.Offset()
	if oc.Y <= 4 {
		cb(occ(cc, 0, -1, 0), 1)
	} else if oc.Y >= 28 {
		cb(occ(cc, 0, 1, 0), 1)
	}
}
<commit_msg>That uses too much memory.<commit_after>package game

import (
	"sync"
	"time"
	"runtime"

	"buildblast/lib/coords"
	"buildblast/lib/mapgen"
)

type ChunkGenerator struct {
	// Chunks are sent to this channel as they are generated
	Generated chan ChunkGenerationResult

	chunks    map[coords.Chunk]ChunkStatus
	mutex     sync.Mutex
	goPool    chan bool
	generator mapgen.Generator
}

type ChunkGenerationResult struct {
	cc    coords.Chunk
	chunk *mapgen.Chunk
}

type ChunkStatus struct {
	generating bool
	generated  bool
	priority   int
}

func NewChunkGenerator(generator mapgen.Generator) *ChunkGenerator {
	cm := new(ChunkGenerator)
	cm.chunks = make(map[coords.Chunk]ChunkStatus, 10)
	cm.Generated = make(chan ChunkGenerationResult, 10)
	// Only cores - 1 chunks can be generated at a time.
	maxActiveThreads := runtime.NumCPU() - 1
	if maxActiveThreads < 1 {
		maxActiveThreads = 1
	}
	cm.goPool = make(chan bool, maxActiveThreads)
	for i := 0; i < maxActiveThreads; i++ {
		cm.goPool <- true
	}
	cm.generator = generator
	return cm
}

// cm.Lock() MUST BE HELD by the caller, or else calling
// this function is unsafe.
func (cm *ChunkGenerator) queue(cc coords.Chunk, priority int) {
	status := cm.chunks[cc]
	if status.generated || status.generating {
		return
	}
	status.priority += priority
	cm.chunks[cc] = status
}

func (cm *ChunkGenerator) QueueChunksNearby(wc coords.World) {
	cm.mutex.Lock()
	defer cm.mutex.Unlock()

	queue := func(cc coords.Chunk, priority int) {
		cm.queue(cc, priority)
	}

	EachChunkNearby(wc, queue)
}

func (cm *ChunkGenerator) Top() (cc coords.Chunk, valid bool) {
	cm.mutex.Lock()
	defer cm.mutex.Unlock()

	highest := -1
	for key, val := range cm.chunks {
		if val.priority > highest && !val.generated && !val.generating {
			highest = val.priority
			cc = key
		}
	}
	if highest != -1 {
		return cc, true
	}
	return cc, false
}

func (cm *ChunkGenerator) Run() {
	for {
		cc, valid := cm.Top()
		if !valid {
			<-time.After(time.Second / 60)
			continue
		}

		cm.mutex.Lock()
		status := cm.chunks[cc]
		status.generating = true
		cm.chunks[cc] = status
		cm.mutex.Unlock()

		go cm.generate(cc)
	}
}

func (cm *ChunkGenerator) generate(cc coords.Chunk) {
	// Attempt to use a slot
	<- cm.goPool
	chunk := cm.generator.Chunk(cc)

	cm.Generated <- ChunkGenerationResult{
		cc:    cc,
		chunk: chunk,
	}

	cm.mutex.Lock()
	status := cm.chunks[cc]
	status.generating = false
	status.generated = true
	cm.chunks[cc] = status
	cm.mutex.Unlock()
	// Free our slot
	cm.goPool <- true
}

func EachChunkNearby(wc coords.World, cb func(cc coords.Chunk, priority int)) {
	occ := func(cc coords.Chunk, x, y, z int) coords.Chunk {
		return coords.Chunk{
			X: cc.X + x,
			Y: cc.Y + y,
			Z: cc.Z + z,
		}
	}

	eachWithin := func(cc coords.Chunk, xdist, ydist, zdist int, cb func(newCC coords.Chunk, dist int)) {
		abs := func(n int) int {
			if n < 0 {
				return -n
			}
			return n
		}
		dist := func(x, y, z int) int {
			return abs(x) + abs(y) + abs(z)
		}

		cb(cc, 0)
		for x := -xdist; x <= xdist; x++ {
			for y := -ydist; y <= ydist; y++ {
				for z := -zdist; z <= zdist; z++ {
					cb(occ(cc, x, y, z), dist(x, y, z))
				}
			}
		}
	}

	cc := wc.Chunk()
	eachWithin(cc, 2, 0, 2, func(newCC coords.Chunk, dist int) {
		// We want to prioritize further away chunks lower, but the
		// priority must be a positive integer.
		cb(newCC, 10-dist)
	})

	oc := wc.Offset()
	if oc.Y <= 4 {
		cb(occ(cc, 0, -1, 0), 1)
	} else if oc.Y >= 28 {
		cb(occ(cc, 0, 1, 0), 1)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

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 e2e

import (
	"math/rand"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/golang/glog"
)

type testSpec struct {
	// The test to run
	test func(c *client.Client) bool
	// The human readable name of this test
	name string
}

type testInfo struct {
	passed bool
	spec   testSpec
}

// Output a summary in the TAP (test anything protocol) format for automated processing.
// See http://testanything.org/ for more info
func outputTAPSummary(infoList []testInfo) {
	glog.Infof("1..%d", len(infoList))
	for i, info := range infoList {
		if info.passed {
			glog.Infof("ok %d - %s", i+1, info.spec.name)
		} else {
			glog.Infof("not ok %d - %s", i+1, info.spec.name)
		}
	}
}

// Fisher-Yates shuffle using the given RNG r
func shuffleTests(tests []testSpec, r *rand.Rand) {
	for i := len(tests) - 1; i > 0; i-- {
		j := r.Intn(i + 1)
		tests[i], tests[j] = tests[j], tests[i]
	}
}

// Run each Go end-to-end-test. This function assumes the
// creation of a test cluster.
func RunE2ETests(authConfig, certDir, host, repoRoot, provider string, orderseed int64, times int, testList []string) {
	testContext = testContextType{authConfig, certDir, host, repoRoot, provider}
	util.ReallyCrash = true
	util.InitLogs()
	defer util.FlushLogs()

	// TODO: Associate a timeout with each test individually.
	go func() {
		defer util.FlushLogs()
		time.Sleep(10 * time.Minute)
		glog.Fatalf("This test has timed out. Cleanup not guaranteed.")
	}()

	tests := []testSpec{
		/*  Disable TestKubernetesROService due to rate limiter issues.
		    TODO: Add this test back when rate limiting is working properly.
				{TestKubernetesROService, "TestKubernetesROService"},
		*/
		{TestKubeletSendsEvent, "TestKubeletSendsEvent"},
		{TestImportantURLs, "TestImportantURLs"},
		{TestPodUpdate, "TestPodUpdate"},
		{TestNetwork, "TestNetwork"},
		{TestClusterDNS, "TestClusterDNS"},
		{TestPodHasServiceEnvVars, "TestPodHasServiceEnvVars"},
		{TestBasic, "TestBasic"},
		{TestPrivate, "TestPrivate"},
		{TestLivenessHttp, "TestLivenessHttp"},
		{TestLivenessExec, "TestLivenessExec"},
	}

	// Check testList for non-existent tests and populate a StringSet with tests to run.
	validTestNames := util.NewStringSet()
	for _, test := range tests {
		validTestNames.Insert(test.name)
	}
	runTestNames := util.NewStringSet()
	for _, testName := range testList {
		if validTestNames.Has(testName) {
			runTestNames.Insert(testName)
		} else {
			glog.Warningf("Requested test %s does not exist", testName)
		}
	}

	// if testList was specified, filter down now before we expand and shuffle
	if len(testList) > 0 {
		newTests := make([]testSpec, 0)
		for i, test := range tests {
			// Check if this test is supposed to run, either if listed explicitly in
			// a --test flag or if no --test flags were supplied.
			if !runTestNames.Has(test.name) {
				glog.Infof("Skipping test %d %s", i+1, test.name)
				continue
			}
			newTests = append(newTests, test)
		}
		tests = newTests
	}
	if times != 1 {
		newTests := make([]testSpec, 0, times*len(tests))
		for i := 0; i < times; i++ {
			newTests = append(newTests, tests...)
		}
		tests = newTests
	}
	if orderseed == 0 {
		// Use low order bits of NanoTime as the default seed. (Using
		// all the bits makes for a long, very similar looking seed
		// between runs.)
		orderseed = time.Now().UnixNano() & (1<<32 - 1)
	}
	// TODO(satnam6502): When the tests are run in parallel we will
	//                   no longer need the shuffle.
	shuffleTests(tests, rand.New(rand.NewSource(orderseed)))
	glog.Infof("Tests shuffled with orderseed %#x\n", orderseed)

	info := []testInfo{}
	passed := true
	for i, test := range tests {
		glog.Infof("Running test %d %s", i+1, test.name)
		// A client is made for each test. This allows us to attribute
		// issues with rate ACLs etc. to a specific test and supports
		// parallel testing.
		testPassed := test.test(loadClientOrDie())
		if !testPassed {
			glog.Infof("        test %d failed", i+1)
			passed = false
		} else {
			glog.Infof("        test %d passed", i+1)
		}
		// TODO: clean up objects created during a test after the test, so cases
		// are independent.
		info = append(info, testInfo{testPassed, test})
	}
	outputTAPSummary(info)
	if !passed {
		glog.Fatalf("At least one test failed")
	} else {
		glog.Infof("All tests pass")
	}
}
<commit_msg>Reinstate ROService test now that rate limit issue has been addressed<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

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 e2e

import (
	"math/rand"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/golang/glog"
)

type testSpec struct {
	// The test to run
	test func(c *client.Client) bool
	// The human readable name of this test
	name string
}

type testInfo struct {
	passed bool
	spec   testSpec
}

// Output a summary in the TAP (test anything protocol) format for automated processing.
// See http://testanything.org/ for more info
func outputTAPSummary(infoList []testInfo) {
	glog.Infof("1..%d", len(infoList))
	for i, info := range infoList {
		if info.passed {
			glog.Infof("ok %d - %s", i+1, info.spec.name)
		} else {
			glog.Infof("not ok %d - %s", i+1, info.spec.name)
		}
	}
}

// Fisher-Yates shuffle using the given RNG r
func shuffleTests(tests []testSpec, r *rand.Rand) {
	for i := len(tests) - 1; i > 0; i-- {
		j := r.Intn(i + 1)
		tests[i], tests[j] = tests[j], tests[i]
	}
}

// Run each Go end-to-end-test. This function assumes the
// creation of a test cluster.
func RunE2ETests(authConfig, certDir, host, repoRoot, provider string, orderseed int64, times int, testList []string) {
	testContext = testContextType{authConfig, certDir, host, repoRoot, provider}
	util.ReallyCrash = true
	util.InitLogs()
	defer util.FlushLogs()

	// TODO: Associate a timeout with each test individually.
	go func() {
		defer util.FlushLogs()
		time.Sleep(10 * time.Minute)
		glog.Fatalf("This test has timed out. Cleanup not guaranteed.")
	}()

	tests := []testSpec{
		{TestKubernetesROService, "TestKubernetesROService"},
		{TestKubeletSendsEvent, "TestKubeletSendsEvent"},
		{TestImportantURLs, "TestImportantURLs"},
		{TestPodUpdate, "TestPodUpdate"},
		{TestNetwork, "TestNetwork"},
		{TestClusterDNS, "TestClusterDNS"},
		{TestPodHasServiceEnvVars, "TestPodHasServiceEnvVars"},
		{TestBasic, "TestBasic"},
		{TestPrivate, "TestPrivate"},
		{TestLivenessHttp, "TestLivenessHttp"},
		{TestLivenessExec, "TestLivenessExec"},
	}

	// Check testList for non-existent tests and populate a StringSet with tests to run.
	validTestNames := util.NewStringSet()
	for _, test := range tests {
		validTestNames.Insert(test.name)
	}
	runTestNames := util.NewStringSet()
	for _, testName := range testList {
		if validTestNames.Has(testName) {
			runTestNames.Insert(testName)
		} else {
			glog.Warningf("Requested test %s does not exist", testName)
		}
	}

	// if testList was specified, filter down now before we expand and shuffle
	if len(testList) > 0 {
		newTests := make([]testSpec, 0)
		for i, test := range tests {
			// Check if this test is supposed to run, either if listed explicitly in
			// a --test flag or if no --test flags were supplied.
			if !runTestNames.Has(test.name) {
				glog.Infof("Skipping test %d %s", i+1, test.name)
				continue
			}
			newTests = append(newTests, test)
		}
		tests = newTests
	}
	if times != 1 {
		newTests := make([]testSpec, 0, times*len(tests))
		for i := 0; i < times; i++ {
			newTests = append(newTests, tests...)
		}
		tests = newTests
	}
	if orderseed == 0 {
		// Use low order bits of NanoTime as the default seed. (Using
		// all the bits makes for a long, very similar looking seed
		// between runs.)
		orderseed = time.Now().UnixNano() & (1<<32 - 1)
	}
	// TODO(satnam6502): When the tests are run in parallel we will
	//                   no longer need the shuffle.
	shuffleTests(tests, rand.New(rand.NewSource(orderseed)))
	glog.Infof("Tests shuffled with orderseed %#x\n", orderseed)

	info := []testInfo{}
	passed := true
	for i, test := range tests {
		glog.Infof("Running test %d %s", i+1, test.name)
		// A client is made for each test. This allows us to attribute
		// issues with rate ACLs etc. to a specific test and supports
		// parallel testing.
		testPassed := test.test(loadClientOrDie())
		if !testPassed {
			glog.Infof("        test %d failed", i+1)
			passed = false
		} else {
			glog.Infof("        test %d passed", i+1)
		}
		// TODO: clean up objects created during a test after the test, so cases
		// are independent.
		info = append(info, testInfo{testPassed, test})
	}
	outputTAPSummary(info)
	if !passed {
		glog.Fatalf("At least one test failed")
	} else {
		glog.Infof("All tests pass")
	}
}
<|endoftext|>


<commit_before>package bitswap

import (
	"sync"
	"time"

	process "gx/ipfs/QmQopLATEYMNg7dVqZRNDfeE2S1yKy8zrRh5xnYiuqeZBn/goprocess"
	procctx "gx/ipfs/QmQopLATEYMNg7dVqZRNDfeE2S1yKy8zrRh5xnYiuqeZBn/goprocess/context"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"

	key "github.com/ipfs/go-ipfs/blocks/key"
	logging "gx/ipfs/QmNQynaz7qfriSUJkiEZUrm2Wen1u3Kj9goZzWtrPyu7XR/go-log"
	peer "gx/ipfs/QmRBqJF7hb8ZSpRcMwUt8hNhydWcxGEhtk81HKq6oUwKvs/go-libp2p-peer"
)

var TaskWorkerCount = 8

func (bs *Bitswap) startWorkers(px process.Process, ctx context.Context) {
	// Start up a worker to handle block requests this node is making
	px.Go(func(px process.Process) {
		bs.providerQueryManager(ctx)
	})

	// Start up workers to handle requests from other nodes for the data on this node
	for i := 0; i < TaskWorkerCount; i++ {
		i := i
		px.Go(func(px process.Process) {
			bs.taskWorker(ctx, i)
		})
	}

	// Start up a worker to manage periodically resending our wantlist out to peers
	px.Go(func(px process.Process) {
		bs.rebroadcastWorker(ctx)
	})

	// Start up a worker to manage sending out provides messages
	px.Go(func(px process.Process) {
		bs.provideCollector(ctx)
	})

	// Spawn up multiple workers to handle incoming blocks
	// consider increasing number if providing blocks bottlenecks
	// file transfers
	px.Go(bs.provideWorker)
}

func (bs *Bitswap) taskWorker(ctx context.Context, id int) {
	idmap := logging.LoggableMap{"ID": id}
	defer log.Info("bitswap task worker shutting down...")
	for {
		log.Event(ctx, "Bitswap.TaskWorker.Loop", idmap)
		select {
		case nextEnvelope := <-bs.engine.Outbox():
			select {
			case envelope, ok := <-nextEnvelope:
				if !ok {
					continue
				}
				log.Event(ctx, "Bitswap.TaskWorker.Work", logging.LoggableMap{
					"ID":     id,
					"Target": envelope.Peer.Pretty(),
					"Block":  envelope.Block.Multihash().B58String(),
				})

				bs.wm.SendBlock(ctx, envelope)
			case <-ctx.Done():
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func (bs *Bitswap) provideWorker(px process.Process) {

	limit := make(chan struct{}, provideWorkerMax)

	limitedGoProvide := func(k key.Key, wid int) {
		defer func() {
			// replace token when done
			<-limit
		}()
		ev := logging.LoggableMap{"ID": wid}

		ctx := procctx.OnClosingContext(px) // derive ctx from px
		defer log.EventBegin(ctx, "Bitswap.ProvideWorker.Work", ev, &k).Done()

		ctx, cancel := context.WithTimeout(ctx, provideTimeout) // timeout ctx
		defer cancel()

		if err := bs.network.Provide(ctx, k); err != nil {
			log.Warning(err)
		}
	}

	// worker spawner, reads from bs.provideKeys until it closes, spawning a
	// _ratelimited_ number of workers to handle each key.
	for wid := 2; ; wid++ {
		ev := logging.LoggableMap{"ID": 1}
		log.Event(procctx.OnClosingContext(px), "Bitswap.ProvideWorker.Loop", ev)

		select {
		case <-px.Closing():
			return
		case k, ok := <-bs.provideKeys:
			if !ok {
				log.Debug("provideKeys channel closed")
				return
			}
			select {
			case <-px.Closing():
				return
			case limit <- struct{}{}:
				go limitedGoProvide(k, wid)
			}
		}
	}
}

func (bs *Bitswap) provideCollector(ctx context.Context) {
	defer close(bs.provideKeys)
	var toProvide []key.Key
	var nextKey key.Key
	var keysOut chan key.Key

	for {
		select {
		case blk, ok := <-bs.newBlocks:
			if !ok {
				log.Debug("newBlocks channel closed")
				return
			}
			if keysOut == nil {
				nextKey = blk.Key()
				keysOut = bs.provideKeys
			} else {
				toProvide = append(toProvide, blk.Key())
			}
		case keysOut <- nextKey:
			if len(toProvide) > 0 {
				nextKey = toProvide[0]
				toProvide = toProvide[1:]
			} else {
				keysOut = nil
			}
		case <-ctx.Done():
			return
		}
	}
}

func (bs *Bitswap) rebroadcastWorker(parent context.Context) {
	ctx, cancel := context.WithCancel(parent)
	defer cancel()

	broadcastSignal := time.NewTicker(rebroadcastDelay.Get())
	defer broadcastSignal.Stop()

	tick := time.NewTicker(10 * time.Second)
	defer tick.Stop()

	for {
		log.Event(ctx, "Bitswap.Rebroadcast.idle")
		select {
		case <-tick.C:
			n := bs.wm.wl.Len()
			if n > 0 {
				log.Debug(n, "keys in bitswap wantlist")
			}
		case <-broadcastSignal.C: // resend unfulfilled wantlist keys
			log.Event(ctx, "Bitswap.Rebroadcast.active")
			entries := bs.wm.wl.Entries()
			if len(entries) == 0 {
				continue
			}
			tctx, cancel := context.WithTimeout(ctx, providerRequestTimeout)
			for _, e := range bs.wm.wl.Entries() {
				e := e
				bs.findKeys <- &blockRequest{
					Key: e.Key,
					Ctx: tctx,
				}
			}
			cancel()
		case <-parent.Done():
			return
		}
	}
}

func (bs *Bitswap) providerQueryManager(ctx context.Context) {
	var activeLk sync.Mutex
	kset := key.NewKeySet()

	for {
		select {
		case e := <-bs.findKeys:
			activeLk.Lock()
			if kset.Has(e.Key) {
				activeLk.Unlock()
				continue
			}
			kset.Add(e.Key)
			activeLk.Unlock()

			go func(e *blockRequest) {
				child, cancel := context.WithTimeout(e.Ctx, providerRequestTimeout)
				defer cancel()
				providers := bs.network.FindProvidersAsync(child, e.Key, maxProvidersPerRequest)
				wg := &sync.WaitGroup{}
				for p := range providers {
					wg.Add(1)
					go func(p peer.ID) {
						defer wg.Done()
						err := bs.network.ConnectTo(child, p)
						if err != nil {
							log.Debug("failed to connect to provider %s: %s", p, err)
						}
					}(p)
				}
				wg.Wait()
				activeLk.Lock()
				kset.Remove(e.Key)
				activeLk.Unlock()
			}(e)

		case <-ctx.Done():
			return
		}
	}
}
<commit_msg>bitswap: search for wantlist providers a little less often<commit_after>package bitswap

import (
	"math/rand"
	"sync"
	"time"

	process "gx/ipfs/QmQopLATEYMNg7dVqZRNDfeE2S1yKy8zrRh5xnYiuqeZBn/goprocess"
	procctx "gx/ipfs/QmQopLATEYMNg7dVqZRNDfeE2S1yKy8zrRh5xnYiuqeZBn/goprocess/context"
	context "gx/ipfs/QmZy2y8t9zQH2a1b8q2ZSLKp17ATuJoCNxxyMFG5qFExpt/go-net/context"

	key "github.com/ipfs/go-ipfs/blocks/key"
	logging "gx/ipfs/QmNQynaz7qfriSUJkiEZUrm2Wen1u3Kj9goZzWtrPyu7XR/go-log"
	peer "gx/ipfs/QmRBqJF7hb8ZSpRcMwUt8hNhydWcxGEhtk81HKq6oUwKvs/go-libp2p-peer"
)

var TaskWorkerCount = 8

func (bs *Bitswap) startWorkers(px process.Process, ctx context.Context) {
	// Start up a worker to handle block requests this node is making
	px.Go(func(px process.Process) {
		bs.providerQueryManager(ctx)
	})

	// Start up workers to handle requests from other nodes for the data on this node
	for i := 0; i < TaskWorkerCount; i++ {
		i := i
		px.Go(func(px process.Process) {
			bs.taskWorker(ctx, i)
		})
	}

	// Start up a worker to manage periodically resending our wantlist out to peers
	px.Go(func(px process.Process) {
		bs.rebroadcastWorker(ctx)
	})

	// Start up a worker to manage sending out provides messages
	px.Go(func(px process.Process) {
		bs.provideCollector(ctx)
	})

	// Spawn up multiple workers to handle incoming blocks
	// consider increasing number if providing blocks bottlenecks
	// file transfers
	px.Go(bs.provideWorker)
}

func (bs *Bitswap) taskWorker(ctx context.Context, id int) {
	idmap := logging.LoggableMap{"ID": id}
	defer log.Info("bitswap task worker shutting down...")
	for {
		log.Event(ctx, "Bitswap.TaskWorker.Loop", idmap)
		select {
		case nextEnvelope := <-bs.engine.Outbox():
			select {
			case envelope, ok := <-nextEnvelope:
				if !ok {
					continue
				}
				log.Event(ctx, "Bitswap.TaskWorker.Work", logging.LoggableMap{
					"ID":     id,
					"Target": envelope.Peer.Pretty(),
					"Block":  envelope.Block.Multihash().B58String(),
				})

				bs.wm.SendBlock(ctx, envelope)
			case <-ctx.Done():
				return
			}
		case <-ctx.Done():
			return
		}
	}
}

func (bs *Bitswap) provideWorker(px process.Process) {

	limit := make(chan struct{}, provideWorkerMax)

	limitedGoProvide := func(k key.Key, wid int) {
		defer func() {
			// replace token when done
			<-limit
		}()
		ev := logging.LoggableMap{"ID": wid}

		ctx := procctx.OnClosingContext(px) // derive ctx from px
		defer log.EventBegin(ctx, "Bitswap.ProvideWorker.Work", ev, &k).Done()

		ctx, cancel := context.WithTimeout(ctx, provideTimeout) // timeout ctx
		defer cancel()

		if err := bs.network.Provide(ctx, k); err != nil {
			log.Warning(err)
		}
	}

	// worker spawner, reads from bs.provideKeys until it closes, spawning a
	// _ratelimited_ number of workers to handle each key.
	for wid := 2; ; wid++ {
		ev := logging.LoggableMap{"ID": 1}
		log.Event(procctx.OnClosingContext(px), "Bitswap.ProvideWorker.Loop", ev)

		select {
		case <-px.Closing():
			return
		case k, ok := <-bs.provideKeys:
			if !ok {
				log.Debug("provideKeys channel closed")
				return
			}
			select {
			case <-px.Closing():
				return
			case limit <- struct{}{}:
				go limitedGoProvide(k, wid)
			}
		}
	}
}

func (bs *Bitswap) provideCollector(ctx context.Context) {
	defer close(bs.provideKeys)
	var toProvide []key.Key
	var nextKey key.Key
	var keysOut chan key.Key

	for {
		select {
		case blk, ok := <-bs.newBlocks:
			if !ok {
				log.Debug("newBlocks channel closed")
				return
			}
			if keysOut == nil {
				nextKey = blk.Key()
				keysOut = bs.provideKeys
			} else {
				toProvide = append(toProvide, blk.Key())
			}
		case keysOut <- nextKey:
			if len(toProvide) > 0 {
				nextKey = toProvide[0]
				toProvide = toProvide[1:]
			} else {
				keysOut = nil
			}
		case <-ctx.Done():
			return
		}
	}
}

func (bs *Bitswap) rebroadcastWorker(parent context.Context) {
	ctx, cancel := context.WithCancel(parent)
	defer cancel()

	broadcastSignal := time.NewTicker(rebroadcastDelay.Get())
	defer broadcastSignal.Stop()

	tick := time.NewTicker(10 * time.Second)
	defer tick.Stop()

	for {
		log.Event(ctx, "Bitswap.Rebroadcast.idle")
		select {
		case <-tick.C:
			n := bs.wm.wl.Len()
			if n > 0 {
				log.Debug(n, "keys in bitswap wantlist")
			}
		case <-broadcastSignal.C: // resend unfulfilled wantlist keys
			log.Event(ctx, "Bitswap.Rebroadcast.active")
			entries := bs.wm.wl.Entries()
			if len(entries) == 0 {
				continue
			}

			// TODO: come up with a better strategy for determining when to search
			// for new providers for blocks.
			i := rand.Intn(len(entries))
			bs.findKeys <- &blockRequest{
				Key: entries[i].Key,
				Ctx: ctx,
			}
		case <-parent.Done():
			return
		}
	}
}

func (bs *Bitswap) providerQueryManager(ctx context.Context) {
	var activeLk sync.Mutex
	kset := key.NewKeySet()

	for {
		select {
		case e := <-bs.findKeys:
			activeLk.Lock()
			if kset.Has(e.Key) {
				activeLk.Unlock()
				continue
			}
			kset.Add(e.Key)
			activeLk.Unlock()

			go func(e *blockRequest) {
				child, cancel := context.WithTimeout(e.Ctx, providerRequestTimeout)
				defer cancel()
				providers := bs.network.FindProvidersAsync(child, e.Key, maxProvidersPerRequest)
				wg := &sync.WaitGroup{}
				for p := range providers {
					wg.Add(1)
					go func(p peer.ID) {
						defer wg.Done()
						err := bs.network.ConnectTo(child, p)
						if err != nil {
							log.Debug("failed to connect to provider %s: %s", p, err)
						}
					}(p)
				}
				wg.Wait()
				activeLk.Lock()
				kset.Remove(e.Key)
				activeLk.Unlock()
			}(e)

		case <-ctx.Done():
			return
		}
	}
}
<|endoftext|>


<commit_before>package expr

import (
	"fmt"
	"math"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/GaryBoone/GoStats/stats"
	"github.com/StackExchange/tcollector/opentsdb"
	"github.com/StackExchange/tsaf/expr/parse"
	"github.com/StackExchange/tsaf/search"
)

const (
	DefDuration = "1h"
	DefPeriod   = "1w"
	DefNum      = 8
)

var Builtins = map[string]parse.Func{
	"avg": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING},
		parse.TYPE_NUMBER,
		[]interface{}{DefDuration},
		Avg,
	},
	"band": {
		[]parse.FuncType{parse.TYPE_QUERY, parse.TYPE_STRING, parse.TYPE_STRING, parse.TYPE_NUMBER},
		parse.TYPE_SERIES,
		[]interface{}{DefDuration, DefPeriod, DefNum},
		nil,
	},
	"dev": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING},
		parse.TYPE_NUMBER,
		[]interface{}{DefDuration},
		Dev,
	},
	"recent": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING},
		parse.TYPE_NUMBER,
		[]interface{}{DefDuration},
		Recent,
	},
	"since": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING},
		parse.TYPE_NUMBER,
		[]interface{}{DefDuration},
		Since,
	},
	"forecastlr": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING, parse.TYPE_NUMBER},
		parse.TYPE_NUMBER,
		nil,
		Forecast_lr,
	},
}

func queryDuration(host, query, duration string, F func(map[string]opentsdb.Point) float64) (r []*Result, err error) {
	q, err := opentsdb.ParseQuery(query)
	if err != nil {
		return
	}
	if err = expandSearch(q); err != nil {
		return
	}
	d, err := ParseDuration(duration)
	if err != nil {
		return
	}
	req := opentsdb.Request{
		Queries: []*opentsdb.Query{q},
		Start:   fmt.Sprintf("%dms-ago", d.Nanoseconds()/1e6),
	}
	s, err := req.Query(host)
	if err != nil {
		return
	}
	for _, res := range s {
		if len(res.DPS) == 0 {
			// do something here?
			continue
		}
		r = append(r, &Result{
			Value: Value(F(res.DPS)),
			Group: res.Tags,
		})
	}
	return
}

func expandSearch(q *opentsdb.Query) error {
	for k, ov := range q.Tags {
		v := ov
		if v == "*" || !strings.Contains(v, "*") || strings.Contains(v, "|") {
			continue
		}
		v = strings.Replace(v, ".", `\.`, -1)
		v = strings.Replace(v, "*", ".*", -1)
		v = "^" + v + "$"
		re := regexp.MustCompile(v)
		var nvs []string
		vs := search.TagValuesByMetricTagKey(q.Metric, k)
		for _, nv := range vs {
			if re.MatchString(nv) {
				nvs = append(nvs, nv)
			}
		}
		if len(nvs) == 0 {
			return fmt.Errorf("expr: no tags matching %s=%s", k, ov)
		}
		q.Tags[k] = strings.Join(nvs, "|")
	}
	return nil
}

func Avg(host, query, duration string) ([]*Result, error) {
	return queryDuration(host, query, duration, avg)
}

// avg returns the mean of x.
func avg(dps map[string]opentsdb.Point) (a float64) {
	for _, v := range dps {
		a += float64(v)
	}
	a /= float64(len(dps))
	return
}

func Dev(host, query, duration string) ([]*Result, error) {
	return queryDuration(host, query, duration, dev)
}

// dev returns the sample standard deviation of x.
func dev(dps map[string]opentsdb.Point) (d float64) {
	a := avg(dps)
	for _, v := range dps {
		d += math.Pow(float64(v)-a, 2)
	}
	// how should we handle len(x) == 1?
	d /= float64(len(dps) - 1)
	return math.Sqrt(d)
}

func Recent(host, query, duration string) ([]*Result, error) {
	return queryDuration(host, query, duration, recent)
}

func recent(dps map[string]opentsdb.Point) (a float64) {
	last := -1
	for k, v := range dps {
		d, err := strconv.Atoi(k)
		if err != nil {
			panic(err)
		}
		if d > last {
			a = float64(v)
		}
	}
	return
}

func Since(host, query, duration string) ([]*Result, error) {
	return queryDuration(host, query, duration, since)
}

func since(dps map[string]opentsdb.Point) (a float64) {
	var last time.Time
	for k := range dps {
		d, err := strconv.ParseInt(k, 10, 64)
		if err != nil {
			panic(err)
		}
		t := time.Unix(d, 0)
		if t.After(last) {
			last = t
		}
	}
	s := time.Since(last)
	return s.Seconds()
}

//forecast_lr Returns the number of seconds until the the series will have value Y according to a
//Linear Regression
func Forecast_lr(host, query, duration string, y float64) (r []*Result, err error) {
	q, err := opentsdb.ParseQuery(query)
	if err != nil {
		return
	}
	expandSearch(q)
	d, err := ParseDuration(duration)
	if err != nil {
		return
	}
	req := opentsdb.Request{
		Queries: []*opentsdb.Query{q},
		Start:   fmt.Sprintf("%dms-ago", d.Nanoseconds()/1e6),
	}
	s, err := req.Query(host)
	if err != nil {
		return
	}
	for _, res := range s {
		if len(res.DPS) == 0 {
			// do something here?
			continue
		}
		r = append(r, &Result{
			Value: Value(forecast_lr(res.DPS, y)),
			Group: res.Tags,
		})
	}
	return
}

func forecast_lr(dps map[string]opentsdb.Point, y_val float64) (a float64) {
	var x []float64
	var y []float64
	for k, v := range dps {
		d, err := strconv.ParseInt(k, 10, 64)
		if err != nil {
			panic(err)
		}
		x = append(x, float64(d))
		y = append(y, float64(v))
	}
	var slope, intercept, _, _, _, _ = stats.LinearRegression(x, y)
	// If the slope is basically 0, return -1 since forecast alerts wouldn't care about things that
	// "already happened". There might be a better way to handle this, but this works for now
	if int64(slope) == 0 {
		return -1
	}
	//Apparently it is okay for slope to be Zero, there is no divide by zero, not sure why
	intercept_time := (y_val - intercept) / slope
	t := time.Unix(int64(intercept_time), 0)
	s := time.Since(t)
	return -s.Seconds()

}
<commit_msg>Add a percentile function that also serves as min/max.<commit_after>package expr

import (
	"fmt"
	"math"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/GaryBoone/GoStats/stats"
	"github.com/StackExchange/tcollector/opentsdb"
	"github.com/StackExchange/tsaf/expr/parse"
	"github.com/StackExchange/tsaf/search"
)

const (
	DefDuration = "1h"
	DefPeriod   = "1w"
	DefNum      = 8
)

var Builtins = map[string]parse.Func{
	"avg": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING},
		parse.TYPE_NUMBER,
		[]interface{}{DefDuration},
		Avg,
	},
	"band": {
		[]parse.FuncType{parse.TYPE_QUERY, parse.TYPE_STRING, parse.TYPE_STRING, parse.TYPE_NUMBER},
		parse.TYPE_SERIES,
		[]interface{}{DefDuration, DefPeriod, DefNum},
		nil,
	},
	"dev": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING},
		parse.TYPE_NUMBER,
		[]interface{}{DefDuration},
		Dev,
	},
	"recent": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING},
		parse.TYPE_NUMBER,
		[]interface{}{DefDuration},
		Recent,
	},
	"since": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING},
		parse.TYPE_NUMBER,
		[]interface{}{DefDuration},
		Since,
	},
	"forecastlr": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING, parse.TYPE_NUMBER},
		parse.TYPE_NUMBER,
		nil,
		Forecast_lr,
	},
	"percentile": {
		[]parse.FuncType{parse.TYPE_SERIES, parse.TYPE_STRING, parse.TYPE_NUMBER},
		parse.TYPE_NUMBER,
		nil,
		Percentile,
	},
}

func queryDuration(host, query, duration string, F func(map[string]opentsdb.Point) float64) (r []*Result, err error) {
	q, err := opentsdb.ParseQuery(query)
	if err != nil {
		return
	}
	if err = expandSearch(q); err != nil {
		return
	}
	d, err := ParseDuration(duration)
	if err != nil {
		return
	}
	req := opentsdb.Request{
		Queries: []*opentsdb.Query{q},
		Start:   fmt.Sprintf("%dms-ago", d.Nanoseconds()/1e6),
	}
	s, err := req.Query(host)
	if err != nil {
		return
	}
	for _, res := range s {
		if len(res.DPS) == 0 {
			// do something here?
			continue
		}
		r = append(r, &Result{
			Value: Value(F(res.DPS)),
			Group: res.Tags,
		})
	}
	return
}

func expandSearch(q *opentsdb.Query) error {
	for k, ov := range q.Tags {
		v := ov
		if v == "*" || !strings.Contains(v, "*") || strings.Contains(v, "|") {
			continue
		}
		v = strings.Replace(v, ".", `\.`, -1)
		v = strings.Replace(v, "*", ".*", -1)
		v = "^" + v + "$"
		re := regexp.MustCompile(v)
		var nvs []string
		vs := search.TagValuesByMetricTagKey(q.Metric, k)
		for _, nv := range vs {
			if re.MatchString(nv) {
				nvs = append(nvs, nv)
			}
		}
		if len(nvs) == 0 {
			return fmt.Errorf("expr: no tags matching %s=%s", k, ov)
		}
		q.Tags[k] = strings.Join(nvs, "|")
	}
	return nil
}

func Avg(host, query, duration string) ([]*Result, error) {
	return queryDuration(host, query, duration, avg)
}

// avg returns the mean of x.
func avg(dps map[string]opentsdb.Point) (a float64) {
	for _, v := range dps {
		a += float64(v)
	}
	a /= float64(len(dps))
	return
}

func Dev(host, query, duration string) ([]*Result, error) {
	return queryDuration(host, query, duration, dev)
}

// dev returns the sample standard deviation of x.
func dev(dps map[string]opentsdb.Point) (d float64) {
	a := avg(dps)
	for _, v := range dps {
		d += math.Pow(float64(v)-a, 2)
	}
	// how should we handle len(x) == 1?
	d /= float64(len(dps) - 1)
	return math.Sqrt(d)
}

func Recent(host, query, duration string) ([]*Result, error) {
	return queryDuration(host, query, duration, recent)
}

func recent(dps map[string]opentsdb.Point) (a float64) {
	last := -1
	for k, v := range dps {
		d, err := strconv.Atoi(k)
		if err != nil {
			panic(err)
		}
		if d > last {
			a = float64(v)
		}
	}
	return
}

func Since(host, query, duration string) ([]*Result, error) {
	return queryDuration(host, query, duration, since)
}

func since(dps map[string]opentsdb.Point) (a float64) {
	var last time.Time
	for k := range dps {
		d, err := strconv.ParseInt(k, 10, 64)
		if err != nil {
			panic(err)
		}
		t := time.Unix(d, 0)
		if t.After(last) {
			last = t
		}
	}
	s := time.Since(last)
	return s.Seconds()
}

//forecast_lr Returns the number of seconds until the the series will have value Y according to a
//Linear Regression
func Forecast_lr(host, query, duration string, y float64) (r []*Result, err error) {
	q, err := opentsdb.ParseQuery(query)
	if err != nil {
		return
	}
	expandSearch(q)
	d, err := ParseDuration(duration)
	if err != nil {
		return
	}
	req := opentsdb.Request{
		Queries: []*opentsdb.Query{q},
		Start:   fmt.Sprintf("%dms-ago", d.Nanoseconds()/1e6),
	}
	s, err := req.Query(host)
	if err != nil {
		return
	}
	for _, res := range s {
		if len(res.DPS) == 0 {
			// do something here?
			continue
		}
		r = append(r, &Result{
			Value: Value(forecast_lr(res.DPS, y)),
			Group: res.Tags,
		})
	}
	return
}

func forecast_lr(dps map[string]opentsdb.Point, y_val float64) (a float64) {
	var x []float64
	var y []float64
	for k, v := range dps {
		d, err := strconv.ParseInt(k, 10, 64)
		if err != nil {
			panic(err)
		}
		x = append(x, float64(d))
		y = append(y, float64(v))
	}
	var slope, intercept, _, _, _, _ = stats.LinearRegression(x, y)
	// If the slope is basically 0, return -1 since forecast alerts wouldn't care about things that
	// "already happened". There might be a better way to handle this, but this works for now
	if int64(slope) == 0 {
		return -1
	}
	//Apparently it is okay for slope to be Zero, there is no divide by zero, not sure why
	intercept_time := (y_val - intercept) / slope
	t := time.Unix(int64(intercept_time), 0)
	s := time.Since(t)
	return -s.Seconds()
}

func Percentile(host, query, duration string, p float64) (r []*Result, err error) {
	if p < 0 || p > 1 {
		return nil, fmt.Errorf("requested percentile must be inclusively between 0 and 1")
	}
	q, err := opentsdb.ParseQuery(query)
	if err != nil {
		return
	}
	expandSearch(q)
	d, err := ParseDuration(duration)
	if err != nil {
		return
	}
	req := opentsdb.Request{
		Queries: []*opentsdb.Query{q},
		Start:   fmt.Sprintf("%dms-ago", d.Nanoseconds()/1e6),
	}
	s, err := req.Query(host)
	if err != nil {
		return
	}
	for _, res := range s {
		if len(res.DPS) == 0 {
			return nil, fmt.Errorf("Can not call percentile on a zero length array")
		}
		r = append(r, &Result{
			Value: Value(percentile(res.DPS, p)),
			Group: res.Tags,
		})
	}
	return
}

//percentile returns the value at the corresponding percentile between 0 and 1. There is no standard
//def of percentile so look the code to see how this one works. This also accepts 0 and 1 as special
//cases and returns min and max respectively
func percentile(dps map[string]opentsdb.Point, p float64) (a float64) {
	var x []float64
	for _, v := range dps {
		x = append(x, float64(v))
	}
	sort.Float64s(x)
	if p == 0 {
		return x[0]
	}
	if p == 1 {
		return x[len(x)-1]
	}
	i := p * float64(len(x)-1)
	i = math.Ceil(i)
	return x[int(i)]
}
<|endoftext|>


<commit_before>package gerrit_test

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"reflect"
	"testing"

	"github.com/andygrunwald/go-gerrit"
)

const (
	// testGerritInstanceURL is a test instance url that won`t be called
	testGerritInstanceURL = "https://go-review.googlesource.com/"
)

var (
	// testMux is the HTTP request multiplexer used with the test server.
	testMux *http.ServeMux

	// testClient is the gerrit client being tested.
	testClient *gerrit.Client

	// testServer is a test HTTP server used to provide mock API responses.
	testServer *httptest.Server
)

type testValues map[string]string

// setup sets up a test HTTP server along with a gerrit.Client that is configured to talk to that test server.
// Tests should register handlers on mux which provide mock responses for the API method being tested.
func setup() {
	// Test server
	testMux = http.NewServeMux()
	testServer = httptest.NewServer(testMux)

	// gerrit client configured to use test server
	testClient, _ = gerrit.NewClient(testServer.URL, nil)
}

// teardown closes the test HTTP server.
func teardown() {
	testServer.Close()
}

func testMethod(t *testing.T, r *http.Request, want string) {
	if got := r.Method; got != want {
		t.Errorf("Request method: %v, want %v", got, want)
	}
}

func testRequestURL(t *testing.T, r *http.Request, want string) {
	if got := r.URL.String(); got != want {
		t.Errorf("Request URL: %v, want %v", got, want)
	}
}

func testFormValues(t *testing.T, r *http.Request, values testValues) {
	want := url.Values{}
	for k, v := range values {
		want.Add(k, v)
	}

	r.ParseForm()
	if got := r.Form; !reflect.DeepEqual(got, want) {
		t.Errorf("Request parameters: %v, want %v", got, want)
	}
}

func TestNewClient_NoGerritInstance(t *testing.T) {
	mockData := []string{"", "://not-existing"}
	for _, data := range mockData {
		c, err := gerrit.NewClient(data, nil)
		if c != nil {
			t.Errorf("NewClient return is not nil. Expected no client. Go %+v", c)
		}
		if err == nil {
			t.Error("No error occured by empty Gerrit Instance. Expected one.")
		}
	}
}

func TestNewClient_Services(t *testing.T) {
	c, err := gerrit.NewClient("https://gerrit-review.googlesource.com/", nil)
	if err != nil {
		t.Errorf("An error occured. Expected nil. Got %+v.", err)
	}

	if c.Authentication == nil {
		t.Error("No AuthenticationService found.")
	}
	if c.Access == nil {
		t.Error("No AccessService found.")
	}
	if c.Accounts == nil {
		t.Error("No AccountsService found.")
	}
	if c.Changes == nil {
		t.Error("No ChangesService found.")
	}
	if c.Config == nil {
		t.Error("No ConfigService found.")
	}
	if c.Groups == nil {
		t.Error("No GroupsService found.")
	}
	if c.Plugins == nil {
		t.Error("No PluginsService found.")
	}
	if c.Projects == nil {
		t.Error("No ProjectsService found.")
	}
}

func TestNewRequest(t *testing.T) {
	c, err := gerrit.NewClient(testGerritInstanceURL, nil)
	if err != nil {
		t.Errorf("An error occured. Expected nil. Got %+v.", err)
	}

	inURL, outURL := "/foo", testGerritInstanceURL+"foo"
	inBody, outBody := &gerrit.PermissionRuleInfo{Action: "ALLOW", Force: true, Min: 0, Max: 0}, `{"action":"ALLOW","force":true,"min":0,"max":0}`+"\n"
	req, _ := c.NewRequest("GET", inURL, inBody)

	// Test that relative URL was expanded
	if got, want := req.URL.String(), outURL; got != want {
		t.Errorf("NewRequest(%q) URL is %v, want %v", inURL, got, want)
	}

	// Test that body was JSON encoded
	body, _ := ioutil.ReadAll(req.Body)
	if got, want := string(body), outBody; got != want {
		t.Errorf("NewRequest Body is %v, want %v", got, want)
	}
}

func testURLParseError(t *testing.T, err error) {
	if err == nil {
		t.Errorf("Expected error to be returned")
	}
	if err, ok := err.(*url.Error); !ok || err.Op != "parse" {
		t.Errorf("Expected URL parse error, got %+v", err)
	}
}

func TestNewRequest_BadURL(t *testing.T) {
	c, err := gerrit.NewClient(testGerritInstanceURL, nil)
	if err != nil {
		t.Errorf("An error occured. Expected nil. Got %+v.", err)
	}
	_, err = c.NewRequest("GET", ":", nil)
	testURLParseError(t, err)
}

// If a nil body is passed to gerrit.NewRequest, make sure that nil is also passed to http.NewRequest.
// In most cases, passing an io.Reader that returns no content is fine,
// since there is no difference between an HTTP request body that is an empty string versus one that is not set at all.
// However in certain cases, intermediate systems may treat these differently resulting in subtle errors.
func TestNewRequest_EmptyBody(t *testing.T) {
	c, err := gerrit.NewClient(testGerritInstanceURL, nil)
	if err != nil {
		t.Errorf("An error occured. Expected nil. Got %+v.", err)
	}
	req, err := c.NewRequest("GET", "/", nil)
	if err != nil {
		t.Fatalf("NewRequest returned unexpected error: %v", err)
	}
	if req.Body != nil {
		t.Fatalf("constructed request contains a non-nil Body")
	}
}

func TestDo(t *testing.T) {
	setup()
	defer teardown()

	type foo struct {
		A string
	}

	testMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if m := "GET"; m != r.Method {
			t.Errorf("Request method = %v, want %v", r.Method, m)
		}
		fmt.Fprint(w, `)]}'`+"\n"+`{"A":"a"}`)
	})

	req, _ := testClient.NewRequest("GET", "/", nil)
	body := new(foo)
	testClient.Do(req, body)

	want := &foo{"a"}
	if !reflect.DeepEqual(body, want) {
		t.Errorf("Response body = %v, want %v", body, want)
	}
}

func TestDo_ioWriter(t *testing.T) {
	setup()
	defer teardown()
	content := `)]}'` + "\n" + `{"A":"a"}`

	testMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if m := "GET"; m != r.Method {
			t.Errorf("Request method = %v, want %v", r.Method, m)
		}
		fmt.Fprint(w, content)
	})

	req, _ := testClient.NewRequest("GET", "/", nil)
	var buf []byte
	actual := bytes.NewBuffer(buf)
	testClient.Do(req, actual)

	expected := []byte(content)
	if !reflect.DeepEqual(actual.Bytes(), expected) {
		t.Errorf("Response body = %v, want %v", actual, string(expected))
	}
}

func TestDo_HTTPError(t *testing.T) {
	setup()
	defer teardown()

	testMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "Bad Request", 400)
	})

	req, _ := testClient.NewRequest("GET", "/", nil)
	_, err := testClient.Do(req, nil)

	if err == nil {
		t.Error("Expected HTTP 400 error.")
	}
}

// Test handling of an error caused by the internal http client's Do() function.
// A redirect loop is pretty unlikely to occur within the Gerrit API, but does allow us to exercise the right code path.
func TestDo_RedirectLoop(t *testing.T) {
	setup()
	defer teardown()

	testMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/", http.StatusFound)
	})

	req, _ := testClient.NewRequest("GET", "/", nil)
	_, err := testClient.Do(req, nil)

	if err == nil {
		t.Error("Expected error to be returned.")
	}
	if err, ok := err.(*url.Error); !ok {
		t.Errorf("Expected a URL error; got %#v.", err)
	}
}

func TestRemoveMagicPrefixLine(t *testing.T) {
	mockData := []struct {
		Current, Expected []byte
	}{
		{[]byte(`{"A":"a"}`), []byte(`{"A":"a"}`)},
		{[]byte(`)]}'` + "\n" + `{"A":"a"}`), []byte(`{"A":"a"}`)},
	}
	for _, mock := range mockData {
		body := gerrit.RemoveMagicPrefixLine(mock.Current)
		if !reflect.DeepEqual(body, mock.Expected) {
			t.Errorf("Response body = %v, want %v", body, mock.Expected)
		}
	}
}

func TestRemoveMagicPrefixLineDoesNothingWithoutPrefix(t *testing.T) {
	mockData := []struct {
		Current, Expected []byte
	}{
		{[]byte(`{"A":"a"}`), []byte(`{"A":"a"}`)},
		{[]byte(`{"A":"a"}`), []byte(`{"A":"a"}`)},
	}
	for _, mock := range mockData {
		body := gerrit.RemoveMagicPrefixLine(mock.Current)
		if !reflect.DeepEqual(body, mock.Expected) {
			t.Errorf("Response body = %v, want %v", body, mock.Expected)
		}
	}
}

func TestErrNoInstanceGiven(t *testing.T) {
	_, err := gerrit.NewClient("", nil)
	if err != gerrit.ErrNoInstanceGiven {
		t.Error("Expected `ErrNoInstanceGiven`")
	}
}
<commit_msg>initial tests<commit_after>package gerrit_test

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"reflect"
	"testing"

	"github.com/andygrunwald/go-gerrit"
)

const (
	// testGerritInstanceURL is a test instance url that won`t be called
	testGerritInstanceURL = "https://go-review.googlesource.com/"
)

var (
	// testMux is the HTTP request multiplexer used with the test server.
	testMux *http.ServeMux

	// testClient is the gerrit client being tested.
	testClient *gerrit.Client

	// testServer is a test HTTP server used to provide mock API responses.
	testServer *httptest.Server
)

type testValues map[string]string

// setup sets up a test HTTP server along with a gerrit.Client that is configured to talk to that test server.
// Tests should register handlers on mux which provide mock responses for the API method being tested.
func setup() {
	// Test server
	testMux = http.NewServeMux()
	testServer = httptest.NewServer(testMux)

	// gerrit client configured to use test server
	testClient, _ = gerrit.NewClient(testServer.URL, nil)
}

// teardown closes the test HTTP server.
func teardown() {
	testServer.Close()
}

func testMethod(t *testing.T, r *http.Request, want string) {
	if got := r.Method; got != want {
		t.Errorf("Request method: %v, want %v", got, want)
	}
}

func testRequestURL(t *testing.T, r *http.Request, want string) {
	if got := r.URL.String(); got != want {
		t.Errorf("Request URL: %v, want %v", got, want)
	}
}

func testFormValues(t *testing.T, r *http.Request, values testValues) {
	want := url.Values{}
	for k, v := range values {
		want.Add(k, v)
	}

	r.ParseForm()
	if got := r.Form; !reflect.DeepEqual(got, want) {
		t.Errorf("Request parameters: %v, want %v", got, want)
	}
}

func TestNewClient_NoGerritInstance(t *testing.T) {
	mockData := []string{"", "://not-existing"}
	for _, data := range mockData {
		c, err := gerrit.NewClient(data, nil)
		if c != nil {
			t.Errorf("NewClient return is not nil. Expected no client. Go %+v", c)
		}
		if err == nil {
			t.Error("No error occured by empty Gerrit Instance. Expected one.")
		}
	}
}

func TestNewClient_Services(t *testing.T) {
	c, err := gerrit.NewClient("https://gerrit-review.googlesource.com/", nil)
	if err != nil {
		t.Errorf("An error occured. Expected nil. Got %+v.", err)
	}

	if c.Authentication == nil {
		t.Error("No AuthenticationService found.")
	}
	if c.Access == nil {
		t.Error("No AccessService found.")
	}
	if c.Accounts == nil {
		t.Error("No AccountsService found.")
	}
	if c.Changes == nil {
		t.Error("No ChangesService found.")
	}
	if c.Config == nil {
		t.Error("No ConfigService found.")
	}
	if c.Groups == nil {
		t.Error("No GroupsService found.")
	}
	if c.Plugins == nil {
		t.Error("No PluginsService found.")
	}
	if c.Projects == nil {
		t.Error("No ProjectsService found.")
	}
}

func TestNewClient_TestErrNoInstanceGiven(t *testing.T) {
	_, err := gerrit.NewClient("", nil)
	if err != gerrit.ErrNoInstanceGiven {
		t.Error("Expected `ErrNoInstanceGiven`")
	}
}

func TestNewClientFromURL_NoCredentials(t *testing.T) {
	client, err := gerrit.NewClientFromURL("http://localhost/", nil)
	if err != nil {
		t.Errorf("Unexpected error: %s", err.Error())
	}
	if client.Authentication.HasAuth() {
		t.Error("Expected HasAuth() to return false")
	}
}

func TestNewClientFromURL_UsernameWithoutPassword(t *testing.T) {
	_, err := gerrit.NewClientFromURL("http://foo@localhost/", nil)
	if err != gerrit.ErrUserProvidedWithoutPassword {
		t.Error("Expected ErrUserProvidedWithoutPassword")
	}
}

func TestNewClientFromURL_AuthenticationFailed(t *testing.T) {
	setup()
	defer teardown()

	serverURL := fmt.Sprintf("http://admin:secret@%s/", testServer.Listener.Addr().String())
	_, err := gerrit.NewClientFromURL(serverURL, nil)

	if err != gerrit.ErrAuthenticationFailed {
		t.Error("Expected ErrAuthenticationFailed")
	}
}

func TestNewRequest(t *testing.T) {
	c, err := gerrit.NewClient(testGerritInstanceURL, nil)
	if err != nil {
		t.Errorf("An error occured. Expected nil. Got %+v.", err)
	}

	inURL, outURL := "/foo", testGerritInstanceURL+"foo"
	inBody, outBody := &gerrit.PermissionRuleInfo{Action: "ALLOW", Force: true, Min: 0, Max: 0}, `{"action":"ALLOW","force":true,"min":0,"max":0}`+"\n"
	req, _ := c.NewRequest("GET", inURL, inBody)

	// Test that relative URL was expanded
	if got, want := req.URL.String(), outURL; got != want {
		t.Errorf("NewRequest(%q) URL is %v, want %v", inURL, got, want)
	}

	// Test that body was JSON encoded
	body, _ := ioutil.ReadAll(req.Body)
	if got, want := string(body), outBody; got != want {
		t.Errorf("NewRequest Body is %v, want %v", got, want)
	}
}

func testURLParseError(t *testing.T, err error) {
	if err == nil {
		t.Errorf("Expected error to be returned")
	}
	if err, ok := err.(*url.Error); !ok || err.Op != "parse" {
		t.Errorf("Expected URL parse error, got %+v", err)
	}
}

func TestNewRequest_BadURL(t *testing.T) {
	c, err := gerrit.NewClient(testGerritInstanceURL, nil)
	if err != nil {
		t.Errorf("An error occured. Expected nil. Got %+v.", err)
	}
	_, err = c.NewRequest("GET", ":", nil)
	testURLParseError(t, err)
}

// If a nil body is passed to gerrit.NewRequest, make sure that nil is also passed to http.NewRequest.
// In most cases, passing an io.Reader that returns no content is fine,
// since there is no difference between an HTTP request body that is an empty string versus one that is not set at all.
// However in certain cases, intermediate systems may treat these differently resulting in subtle errors.
func TestNewRequest_EmptyBody(t *testing.T) {
	c, err := gerrit.NewClient(testGerritInstanceURL, nil)
	if err != nil {
		t.Errorf("An error occured. Expected nil. Got %+v.", err)
	}
	req, err := c.NewRequest("GET", "/", nil)
	if err != nil {
		t.Fatalf("NewRequest returned unexpected error: %v", err)
	}
	if req.Body != nil {
		t.Fatalf("constructed request contains a non-nil Body")
	}
}

func TestDo(t *testing.T) {
	setup()
	defer teardown()

	type foo struct {
		A string
	}

	testMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if m := "GET"; m != r.Method {
			t.Errorf("Request method = %v, want %v", r.Method, m)
		}
		fmt.Fprint(w, `)]}'`+"\n"+`{"A":"a"}`)
	})

	req, _ := testClient.NewRequest("GET", "/", nil)
	body := new(foo)
	testClient.Do(req, body)

	want := &foo{"a"}
	if !reflect.DeepEqual(body, want) {
		t.Errorf("Response body = %v, want %v", body, want)
	}
}

func TestDo_ioWriter(t *testing.T) {
	setup()
	defer teardown()
	content := `)]}'` + "\n" + `{"A":"a"}`

	testMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if m := "GET"; m != r.Method {
			t.Errorf("Request method = %v, want %v", r.Method, m)
		}
		fmt.Fprint(w, content)
	})

	req, _ := testClient.NewRequest("GET", "/", nil)
	var buf []byte
	actual := bytes.NewBuffer(buf)
	testClient.Do(req, actual)

	expected := []byte(content)
	if !reflect.DeepEqual(actual.Bytes(), expected) {
		t.Errorf("Response body = %v, want %v", actual, string(expected))
	}
}

func TestDo_HTTPError(t *testing.T) {
	setup()
	defer teardown()

	testMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "Bad Request", 400)
	})

	req, _ := testClient.NewRequest("GET", "/", nil)
	_, err := testClient.Do(req, nil)

	if err == nil {
		t.Error("Expected HTTP 400 error.")
	}
}

// Test handling of an error caused by the internal http client's Do() function.
// A redirect loop is pretty unlikely to occur within the Gerrit API, but does allow us to exercise the right code path.
func TestDo_RedirectLoop(t *testing.T) {
	setup()
	defer teardown()

	testMux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.Redirect(w, r, "/", http.StatusFound)
	})

	req, _ := testClient.NewRequest("GET", "/", nil)
	_, err := testClient.Do(req, nil)

	if err == nil {
		t.Error("Expected error to be returned.")
	}
	if err, ok := err.(*url.Error); !ok {
		t.Errorf("Expected a URL error; got %#v.", err)
	}
}

func TestRemoveMagicPrefixLine(t *testing.T) {
	mockData := []struct {
		Current, Expected []byte
	}{
		{[]byte(`{"A":"a"}`), []byte(`{"A":"a"}`)},
		{[]byte(`)]}'` + "\n" + `{"A":"a"}`), []byte(`{"A":"a"}`)},
	}
	for _, mock := range mockData {
		body := gerrit.RemoveMagicPrefixLine(mock.Current)
		if !reflect.DeepEqual(body, mock.Expected) {
			t.Errorf("Response body = %v, want %v", body, mock.Expected)
		}
	}
}

func TestRemoveMagicPrefixLineDoesNothingWithoutPrefix(t *testing.T) {
	mockData := []struct {
		Current, Expected []byte
	}{
		{[]byte(`{"A":"a"}`), []byte(`{"A":"a"}`)},
		{[]byte(`{"A":"a"}`), []byte(`{"A":"a"}`)},
	}
	for _, mock := range mockData {
		body := gerrit.RemoveMagicPrefixLine(mock.Current)
		if !reflect.DeepEqual(body, mock.Expected) {
			t.Errorf("Response body = %v, want %v", body, mock.Expected)
		}
	}
}

<|endoftext|>


<commit_before>package worker

import (
	"encoding/json"
	"fmt"
	"github.com/go-redis/redis"
	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/clients/metrics"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/engine"
	flotillaLog "github.com/stitchfix/flotilla-os/log"
	"github.com/stitchfix/flotilla-os/queue"
	"github.com/stitchfix/flotilla-os/state"
	"gopkg.in/tomb.v2"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"time"
)

type statusWorker struct {
	sm                       state.Manager
	ee                       engine.Engine
	conf                     config.Config
	log                      flotillaLog.Logger
	pollInterval             time.Duration
	t                        tomb.Tomb
	engine                   *string
	redisClient              *redis.Client
	workerId                 string
	exceptionExtractorClient *http.Client
	exceptionExtractorUrl    string
}

func (sw *statusWorker) Initialize(conf config.Config, sm state.Manager, eksEngine engine.Engine, emrEngine engine.Engine, log flotillaLog.Logger, pollInterval time.Duration, qm queue.Manager) error {
	sw.pollInterval = pollInterval
	sw.conf = conf
	sw.sm = sm
	sw.ee = eksEngine
	sw.log = log
	sw.workerId = fmt.Sprintf("workerid:%d", rand.Int())
	sw.engine = &state.EKSEngine
	if sw.conf.IsSet("eks.exception_extractor_url") {
		sw.exceptionExtractorClient = &http.Client{
			Timeout: time.Second * 5,
		}
		sw.exceptionExtractorUrl = sw.conf.GetString("eks.exception_extractor_url")
	}
	sw.setupRedisClient(conf)
	_ = sw.log.Log("message", "initialized a status worker")
	return nil
}

func (sw *statusWorker) setupRedisClient(conf config.Config) {
	if *sw.engine == state.EKSEngine {
		sw.redisClient = redis.NewClient(&redis.Options{Addr: conf.GetString("redis_address"), DB: conf.GetInt("redis_db")})
	}
}

func (sw *statusWorker) GetTomb() *tomb.Tomb {
	return &sw.t
}

//
// Run updates status of tasks
//
func (sw *statusWorker) Run() error {
	for {
		select {
		case <-sw.t.Dying():
			sw.log.Log("message", "A status worker was terminated")
			return nil
		default:
			if *sw.engine == state.EKSEngine {
				sw.runOnceEKS()
				time.Sleep(sw.pollInterval)
			}
		}
	}
}

func (sw *statusWorker) runOnceEKS() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -30).Format(time.RFC3339),
		},
		"task_type": {state.DefaultTaskType},
		"status":    {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	sw.processEKSRuns(runs)
}

func (sw *statusWorker) processEKSRuns(runs []state.Run) {
	var lockedRuns []state.Run
	for _, run := range runs {
		duration := time.Duration(45) * time.Second
		lock := sw.acquireLock(run, "status", duration)
		if lock {
			lockedRuns = append(lockedRuns, run)
		}
	}
	_ = metrics.Increment(metrics.StatusWorkerLockedRuns, []string{sw.workerId}, float64(len(lockedRuns)))
	for _, run := range lockedRuns {
		start := time.Now()
		_ = sw.log.Log("message", "launching go process eks run", "run", run.RunID)
		go sw.processEKSRun(run)
		_ = metrics.Timing(metrics.StatusWorkerProcessEKSRun, time.Since(start), []string{sw.workerId}, 1)
	}
}
func (sw *statusWorker) acquireLock(run state.Run, purpose string, expiration time.Duration) bool {
	start := time.Now()
	key := fmt.Sprintf("%s-%s", run.RunID, purpose)
	ttl, err := sw.redisClient.TTL(key).Result()
	if err == nil && ttl.Nanoseconds() < 0 {
		_, err = sw.redisClient.Del(key).Result()
	}
	set, err := sw.redisClient.SetNX(key, sw.workerId, expiration).Result()
	if err != nil {
		_ = sw.log.Log("message", "unable to set lock", "error", fmt.Sprintf("%+v", err))
		return true
	}
	_ = metrics.Timing(metrics.StatusWorkerAcquireLock, time.Since(start), []string{sw.workerId}, 1)
	return set
}

func (sw *statusWorker) processEKSRun(run state.Run) {
	_ = sw.log.Log("message", "process eks run", "run", run.RunID)
	reloadRun, err := sw.sm.GetRun(run.RunID)
	if err == nil && reloadRun.Status == state.StatusStopped {
		// Run was updated by another worker process.
		return
	}
	start := time.Now()
	updatedRunWithMetrics, _ := sw.ee.FetchPodMetrics(run)
	_ = metrics.Timing(metrics.StatusWorkerFetchPodMetrics, time.Since(start), []string{sw.workerId}, 1)

	start = time.Now()
	updatedRun, err := sw.ee.FetchUpdateStatus(updatedRunWithMetrics)
	if err != nil {
		_ = sw.log.Log("message", "fetch update status", "run", run.RunID, "error", fmt.Sprintf("%+v", err))
	}
	_ = metrics.Timing(metrics.StatusWorkerFetchUpdateStatus, time.Since(start), []string{sw.workerId}, 1)

	if err == nil {
		subRuns, err := sw.sm.ListRuns(1000, 0, "status", "desc", nil, map[string]string{"PARENT_FLOTILLA_RUN_ID": run.RunID}, state.Engines)
		if err == nil && subRuns.Total > 0 {
			var spawnedRuns state.SpawnedRuns
			for _, subRun := range subRuns.Runs {
				spawnedRuns = append(spawnedRuns, state.SpawnedRun{RunID: subRun.RunID})
			}
			updatedRun.SpawnedRuns = &spawnedRuns
		}
	}
	if err != nil {
		message := fmt.Sprintf("%+v", err)
		minutesInQueue := time.Now().Sub(*run.QueuedAt).Minutes()
		if strings.Contains(message, "not found") && minutesInQueue > float64(30) {
			stoppedAt := time.Now()
			reason := "Job either timed out or not found on the EKS cluster."
			updatedRun.Status = state.StatusStopped
			updatedRun.FinishedAt = &stoppedAt
			updatedRun.ExitReason = &reason
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}

	} else {
		if run.Status != updatedRun.Status && (updatedRun.PodName == run.PodName) {
			sw.logStatusUpdate(updatedRun)
			if updatedRun.ExitCode != nil {
				go sw.cleanupRun(run.RunID)
			}
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			if err != nil {
				_ = sw.log.Log("message", "unable to save eks runs", "error", fmt.Sprintf("%+v", err))
			}

			if updatedRun.Status == state.StatusStopped {
				//TODO - move to a separate worker.
				//_ = sw.eksEngine.Terminate(run)
			}
		} else {
			if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
				updatedRun.MaxCpuUsed != run.MaxCpuUsed ||
				updatedRun.Cpu != run.Cpu ||
				updatedRun.PodName != run.PodName ||
				updatedRun.Memory != run.Memory ||
				updatedRun.PodEvents != run.PodEvents ||
				updatedRun.SpawnedRuns != run.SpawnedRuns {
				_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			}
		}
	}
}

func (sw *statusWorker) cleanupRun(runID string) {
	//Logs maybe delayed before being persisted to S3.
	time.Sleep(120 * time.Second)
	run, err := sw.sm.GetRun(runID)
	if err == nil {
		//Delete run from Kubernetes
		_ = sw.ee.Terminate(run)
	}
}

func (sw *statusWorker) extractExceptions(runID string) {
	//Logs maybe delayed before being persisted to S3.
	time.Sleep(60 * time.Second)
	run, err := sw.sm.GetRun(runID)
	if err == nil {
		jobUrl := fmt.Sprintf("%s/extract/%s", sw.exceptionExtractorUrl, run.RunID)
		res, err := sw.exceptionExtractorClient.Get(jobUrl)
		if err == nil && res != nil && res.Body != nil {
			body, err := ioutil.ReadAll(res.Body)
			if body != nil {
				defer res.Body.Close()
				runExceptions := state.RunExceptions{}
				err = json.Unmarshal(body, &runExceptions)
				if err == nil {
					run.RunExceptions = &runExceptions
				}
			}
			_, _ = sw.sm.UpdateRun(run.RunID, run)
		}
	}
}

func (sw *statusWorker) processEKSRunMetrics(run state.Run) {
	updatedRun, err := sw.ee.FetchPodMetrics(run)
	if err == nil {
		if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
			updatedRun.MaxCpuUsed != run.MaxCpuUsed {
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}
	}
}

func (sw *statusWorker) logStatusUpdate(update state.Run) {
	var err error
	var startedAt, finishedAt time.Time
	var env state.EnvList
	var command string

	if update.StartedAt != nil {
		startedAt = *update.StartedAt
	}

	if update.FinishedAt != nil {
		finishedAt = *update.FinishedAt
	}

	if update.Env != nil {
		env = *update.Env
	}

	if update.Command != nil {
		command = *update.Command
	}

	if update.ExitCode != nil {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"command", command,
			"exit_code", *update.ExitCode,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	} else {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"command", command,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	}

	if err != nil {
		sw.log.Log("message", "Failed to emit status event", "run_id", update.RunID, "error", err.Error())
	}
}

func (sw *statusWorker) findRun(taskArn string) (state.Run, error) {
	var engines []string
	if sw.engine != nil {
		engines = []string{*sw.engine}
	} else {
		engines = nil
	}

	runs, err := sw.sm.ListRuns(1, 0, "started_at", "asc", map[string][]string{
		"task_arn": {taskArn},
	}, nil, engines)
	if err != nil {
		return state.Run{}, errors.Wrapf(err, "problem finding run by task arn [%s]", taskArn)
	}
	if runs.Total > 0 && len(runs.Runs) > 0 {
		return runs.Runs[0], nil
	}
	return state.Run{}, errors.Errorf("no run found for [%s]", taskArn)
}
<commit_msg>adding timeout handling for emr/spark jobs<commit_after>package worker

import (
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/go-redis/redis"
	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/clients/metrics"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/engine"
	flotillaLog "github.com/stitchfix/flotilla-os/log"
	"github.com/stitchfix/flotilla-os/queue"
	"github.com/stitchfix/flotilla-os/state"
	"gopkg.in/tomb.v2"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"time"
)

type statusWorker struct {
	sm                       state.Manager
	ee                       engine.Engine
	conf                     config.Config
	log                      flotillaLog.Logger
	pollInterval             time.Duration
	t                        tomb.Tomb
	engine                   *string
	redisClient              *redis.Client
	workerId                 string
	exceptionExtractorClient *http.Client
	exceptionExtractorUrl    string
	emrEngine                engine.Engine
}

func (sw *statusWorker) Initialize(conf config.Config, sm state.Manager, eksEngine engine.Engine, emrEngine engine.Engine, log flotillaLog.Logger, pollInterval time.Duration, qm queue.Manager) error {
	sw.pollInterval = pollInterval
	sw.conf = conf
	sw.sm = sm
	sw.ee = eksEngine
	sw.log = log
	sw.workerId = fmt.Sprintf("workerid:%d", rand.Int())
	sw.engine = &state.EKSEngine
	sw.emrEngine = emrEngine
	if sw.conf.IsSet("eks.exception_extractor_url") {
		sw.exceptionExtractorClient = &http.Client{
			Timeout: time.Second * 5,
		}
		sw.exceptionExtractorUrl = sw.conf.GetString("eks.exception_extractor_url")
	}
	sw.setupRedisClient(conf)
	_ = sw.log.Log("message", "initialized a status worker")
	return nil
}

func (sw *statusWorker) setupRedisClient(conf config.Config) {
	if *sw.engine == state.EKSEngine {
		sw.redisClient = redis.NewClient(&redis.Options{Addr: conf.GetString("redis_address"), DB: conf.GetInt("redis_db")})
	}
}

func (sw *statusWorker) GetTomb() *tomb.Tomb {
	return &sw.t
}

//
// Run updates status of tasks
//
func (sw *statusWorker) Run() error {
	for {
		select {
		case <-sw.t.Dying():
			sw.log.Log("message", "A status worker was terminated")
			return nil
		default:
			if *sw.engine == state.EKSEngine {
				sw.runOnceEKS()
				sw.runOnceEMR()
				time.Sleep(sw.pollInterval)
			}
		}
	}
}

func (sw *statusWorker) runOnceEMR() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -30).Format(time.RFC3339),
		},
		"task_type": {state.DefaultTaskType},
		"status":    {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSSparkEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	sw.processEMRRuns(runs)
}

func (sw *statusWorker) processEMRRuns(runs []state.Run) {
	for _, run := range runs {
		if run.QueuedAt != nil && run.ActiveDeadlineSeconds != nil {
			runningDuration := time.Now().Sub(*run.QueuedAt)
			if int64(runningDuration.Seconds()) > *run.ActiveDeadlineSeconds {
				err := sw.emrEngine.Terminate(run)
				if err == nil {
					exitCode := int64(1)
					finishedAt := time.Now()
					_, err = sw.sm.UpdateRun(run.RunID, state.Run{
						Status:     state.StatusStopped,
						ExitReason: aws.String(fmt.Sprintf("Job exceeded specified timeout of %v seconds", *run.ActiveDeadlineSeconds)),
						ExitCode:   &exitCode,
						FinishedAt: &finishedAt,
					})
				}
			}
		}
	}
}

func (sw *statusWorker) runOnceEKS() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -30).Format(time.RFC3339),
		},
		"task_type": {state.DefaultTaskType},
		"status":    {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	sw.processEKSRuns(runs)
}

func (sw *statusWorker) processEKSRuns(runs []state.Run) {
	var lockedRuns []state.Run
	for _, run := range runs {
		duration := time.Duration(45) * time.Second
		lock := sw.acquireLock(run, "status", duration)
		if lock {
			lockedRuns = append(lockedRuns, run)
		}
	}
	_ = metrics.Increment(metrics.StatusWorkerLockedRuns, []string{sw.workerId}, float64(len(lockedRuns)))
	for _, run := range lockedRuns {
		start := time.Now()
		_ = sw.log.Log("message", "launching go process eks run", "run", run.RunID)
		go sw.processEKSRun(run)
		_ = metrics.Timing(metrics.StatusWorkerProcessEKSRun, time.Since(start), []string{sw.workerId}, 1)
	}
}
func (sw *statusWorker) acquireLock(run state.Run, purpose string, expiration time.Duration) bool {
	start := time.Now()
	key := fmt.Sprintf("%s-%s", run.RunID, purpose)
	ttl, err := sw.redisClient.TTL(key).Result()
	if err == nil && ttl.Nanoseconds() < 0 {
		_, err = sw.redisClient.Del(key).Result()
	}
	set, err := sw.redisClient.SetNX(key, sw.workerId, expiration).Result()
	if err != nil {
		_ = sw.log.Log("message", "unable to set lock", "error", fmt.Sprintf("%+v", err))
		return true
	}
	_ = metrics.Timing(metrics.StatusWorkerAcquireLock, time.Since(start), []string{sw.workerId}, 1)
	return set
}

func (sw *statusWorker) processEKSRun(run state.Run) {
	_ = sw.log.Log("message", "process eks run", "run", run.RunID)
	reloadRun, err := sw.sm.GetRun(run.RunID)
	if err == nil && reloadRun.Status == state.StatusStopped {
		// Run was updated by another worker process.
		return
	}
	start := time.Now()
	updatedRunWithMetrics, _ := sw.ee.FetchPodMetrics(run)
	_ = metrics.Timing(metrics.StatusWorkerFetchPodMetrics, time.Since(start), []string{sw.workerId}, 1)

	start = time.Now()
	updatedRun, err := sw.ee.FetchUpdateStatus(updatedRunWithMetrics)
	if err != nil {
		_ = sw.log.Log("message", "fetch update status", "run", run.RunID, "error", fmt.Sprintf("%+v", err))
	}
	_ = metrics.Timing(metrics.StatusWorkerFetchUpdateStatus, time.Since(start), []string{sw.workerId}, 1)

	if err == nil {
		subRuns, err := sw.sm.ListRuns(1000, 0, "status", "desc", nil, map[string]string{"PARENT_FLOTILLA_RUN_ID": run.RunID}, state.Engines)
		if err == nil && subRuns.Total > 0 {
			var spawnedRuns state.SpawnedRuns
			for _, subRun := range subRuns.Runs {
				spawnedRuns = append(spawnedRuns, state.SpawnedRun{RunID: subRun.RunID})
			}
			updatedRun.SpawnedRuns = &spawnedRuns
		}
	}
	if err != nil {
		message := fmt.Sprintf("%+v", err)
		minutesInQueue := time.Now().Sub(*run.QueuedAt).Minutes()
		if strings.Contains(message, "not found") && minutesInQueue > float64(30) {
			stoppedAt := time.Now()
			reason := "Job either timed out or not found on the EKS cluster."
			updatedRun.Status = state.StatusStopped
			updatedRun.FinishedAt = &stoppedAt
			updatedRun.ExitReason = &reason
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}

	} else {
		if run.Status != updatedRun.Status && (updatedRun.PodName == run.PodName) {
			sw.logStatusUpdate(updatedRun)
			if updatedRun.ExitCode != nil {
				go sw.cleanupRun(run.RunID)
			}
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			if err != nil {
				_ = sw.log.Log("message", "unable to save eks runs", "error", fmt.Sprintf("%+v", err))
			}

			if updatedRun.Status == state.StatusStopped {
				//TODO - move to a separate worker.
				//_ = sw.eksEngine.Terminate(run)
			}
		} else {
			if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
				updatedRun.MaxCpuUsed != run.MaxCpuUsed ||
				updatedRun.Cpu != run.Cpu ||
				updatedRun.PodName != run.PodName ||
				updatedRun.Memory != run.Memory ||
				updatedRun.PodEvents != run.PodEvents ||
				updatedRun.SpawnedRuns != run.SpawnedRuns {
				_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			}
		}
	}
}

func (sw *statusWorker) cleanupRun(runID string) {
	//Logs maybe delayed before being persisted to S3.
	time.Sleep(120 * time.Second)
	run, err := sw.sm.GetRun(runID)
	if err == nil {
		//Delete run from Kubernetes
		_ = sw.ee.Terminate(run)
	}
}

func (sw *statusWorker) extractExceptions(runID string) {
	//Logs maybe delayed before being persisted to S3.
	time.Sleep(60 * time.Second)
	run, err := sw.sm.GetRun(runID)
	if err == nil {
		jobUrl := fmt.Sprintf("%s/extract/%s", sw.exceptionExtractorUrl, run.RunID)
		res, err := sw.exceptionExtractorClient.Get(jobUrl)
		if err == nil && res != nil && res.Body != nil {
			body, err := ioutil.ReadAll(res.Body)
			if body != nil {
				defer res.Body.Close()
				runExceptions := state.RunExceptions{}
				err = json.Unmarshal(body, &runExceptions)
				if err == nil {
					run.RunExceptions = &runExceptions
				}
			}
			_, _ = sw.sm.UpdateRun(run.RunID, run)
		}
	}
}

func (sw *statusWorker) processEKSRunMetrics(run state.Run) {
	updatedRun, err := sw.ee.FetchPodMetrics(run)
	if err == nil {
		if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
			updatedRun.MaxCpuUsed != run.MaxCpuUsed {
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}
	}
}

func (sw *statusWorker) logStatusUpdate(update state.Run) {
	var err error
	var startedAt, finishedAt time.Time
	var env state.EnvList
	var command string

	if update.StartedAt != nil {
		startedAt = *update.StartedAt
	}

	if update.FinishedAt != nil {
		finishedAt = *update.FinishedAt
	}

	if update.Env != nil {
		env = *update.Env
	}

	if update.Command != nil {
		command = *update.Command
	}

	if update.ExitCode != nil {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"command", command,
			"exit_code", *update.ExitCode,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	} else {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"command", command,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	}

	if err != nil {
		sw.log.Log("message", "Failed to emit status event", "run_id", update.RunID, "error", err.Error())
	}
}

func (sw *statusWorker) findRun(taskArn string) (state.Run, error) {
	var engines []string
	if sw.engine != nil {
		engines = []string{*sw.engine}
	} else {
		engines = nil
	}

	runs, err := sw.sm.ListRuns(1, 0, "started_at", "asc", map[string][]string{
		"task_arn": {taskArn},
	}, nil, engines)
	if err != nil {
		return state.Run{}, errors.Wrapf(err, "problem finding run by task arn [%s]", taskArn)
	}
	if runs.Total > 0 && len(runs.Runs) > 0 {
		return runs.Runs[0], nil
	}
	return state.Run{}, errors.Errorf("no run found for [%s]", taskArn)
}
<|endoftext|>