<commit_before>package yasha

import (
	"compress/bzip2"
	"fmt"
	"io/ioutil"
	"net/http"
	"testing"
	"time"

	"github.com/dotabuff/yasha/dota"
	"github.com/stretchr/testify/assert"
)

type testCase struct {
	matchId int64
	url     string

	expectLastChatMessage string
}

// Secret vs Cloud 9 played on patch 6.83c
func TestEsportsPatch683b(t *testing.T) {
	c := &testCase{
		matchId: 1405240741,
		url:     "http://replay135.valve.net/570/1405240741_220241732.dem.bz2",
		expectLastChatMessage: "Gg",
	}

	testReplayCase(t, c)
}

// Navi vs Basically Unknown, played on patch 6.84p0
func TestEsportsPatch684p0(t *testing.T) {
	c := &testCase{
		matchId: 1450235906,
		url:     "http://replay136.valve.net/570/1450235906_1463120933.dem.bz2",
		expectLastChatMessage: "gg",
	}

	testReplayCase(t, c)
}

// No Respeta Funadores vs Who Needs Skill, played on patch 6.84p1
func TestEsportsPatch684p1(t *testing.T) {
	c := &testCase{
		matchId: 1458895412,
		url:     "http://replay123.valve.net/570/1458895412_140022944.dem.bz2",
		expectLastChatMessage: "gg",
	}

	testReplayCase(t, c)
}

// Manually scrutinised match, played on patch 6.84p1
func TestPublicMatchPatch684p1(t *testing.T) {
	assert := assert.New(t)

	data, err := getReplayData(1456774107, "http://s.tsai.co/replays/1456774107.dem")
	if err != nil {
		t.Fatalf("unable to get replay: %s", err)
	}

	parser := NewParser(data)
	parser.OnSayText2 = func(n int, o *dota.CUserMsg_SayText2) {
	}

	earthshakerDeaths := 0
	spiritBreakerDeaths := 0
	parser.OnCombatLog = func(entry CombatLogEntry) {
		// t.Logf("OnCombatLog: %s: %+v", reflect.TypeOf(entry), entry)
		switch log := entry.(type) {
		case *CombatLogDeath:
			if log.Target == "npc_dota_hero_earthshaker" {
				earthshakerDeaths++
			}
			if log.Target == "npc_dota_hero_spirit_breaker" {
				spiritBreakerDeaths++
			}
		}
	}

	var now time.Duration
	var gameTime, preGameStarttime float64
	parser.OnEntityPreserved = func(pe *PacketEntity) {
		if pe.Name == "DT_DOTAGamerulesProxy" {
			gameTime = pe.Values["DT_DOTAGamerules.m_fGameTime"].(float64)
			preGameStarttime = pe.Values["DT_DOTAGamerules.m_flPreGameStartTime"].(float64)
			now = time.Duration(gameTime-preGameStarttime) * time.Second
		}
	}

	// entindex:3 order_type:1 units:349 position:<x:6953.3125 y:6920.8438 z:384 > queue:false
	unitOrderCount := 0
	unitOrderQueuedCount := 0
	specificUnitOrder := false
	parser.OnSpectatorPlayerUnitOrders = func(n int, o *dota.CDOTAUserMsg_SpectatorPlayerUnitOrders) {
		unitOrderCount++
		if *o.Queue == true {
			unitOrderQueuedCount++
		}
		if *o.Entindex == 3 && *o.OrderType == 1 && o.Units[0] == 349 && *o.Queue == false &&
			*o.Position.X == 6953.3125 && *o.Position.Y == 6920.8438 && *o.Position.Y == 384.0 {
			specificUnitOrder = true
		}
	}

	chatWheelMessagesCount := 0
	parser.OnChatWheel = func(n int, o *dota.CDOTAUserMsg_ChatWheel) {
		chatWheelMessagesCount++
	}

	parser.Parse()

	assert.Equal(8, earthshakerDeaths)
	assert.Equal(11, spiritBreakerDeaths)          // not actually right but verified in replay
	assert.Equal(55316, unitOrderCount)            // regression test
	assert.Equal(102, unitOrderQueuedCount)        // regression test
	assert.Equal(int64(2585000000000), int64(now)) // regression test
	assert.Equal(0, chatWheelMessagesCount)        // regression test
}

func testReplayCase(t *testing.T, c *testCase) {
	assert := assert.New(t)

	data, err := getReplayData(c.matchId, c.url)
	if err != nil {
		t.Fatalf("unable to get replay: %s", err)
	}

	lastChatMessage := ""

	parser := NewParser(data)
	parser.OnSayText2 = func(n int, o *dota.CUserMsg_SayText2) {
		//t.Logf("OnSayText2: %+v", o)
		lastChatMessage = o.GetText()
	}

	parser.OnChatEvent = func(n int, o *dota.CDOTAUserMsg_ChatEvent) {
		//t.Logf("OnChatEvent: %+v", o)
	}
	parser.Parse()

	assert.Equal(c.expectLastChatMessage, lastChatMessage)
}

func getReplayData(matchId int64, url string) ([]byte, error) {
	path := fmt.Sprintf("replays/%d.dem", matchId)
	if data, err := ioutil.ReadFile(path); err == nil {
		fmt.Printf("read replay %d from %s\n", matchId, path)
		return data, nil
	}

	fmt.Printf("downloading replay %d from %s...\n", matchId, url)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Return an error if we don't get a 200
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("invalid status %d", resp.StatusCode)
	}

	var data []byte
	if url[len(url)-3:] == "bz2" {
		data, err = ioutil.ReadAll(bzip2.NewReader(resp.Body))
	} else {
		data, err = ioutil.ReadAll(resp.Body)
	}

	if err != nil {
		return nil, err
	}

	if err := ioutil.WriteFile(path, data, 0644); err != nil {
		return nil, err
	}

	fmt.Printf("downloaded replay %d from %s to %s\n", matchId, url, path)

	return data, nil
}
<commit_msg>Improve tests<commit_after>package yasha

import (
	"compress/bzip2"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/dotabuff/yasha/dota"
	"github.com/stretchr/testify/assert"
)

type testCase struct {
	matchId int64
	url     string

	expectLastChatMessage string
	expectHeroKillCount   map[string]int
	expectHeroDeathCount  map[string]int
}

// Esports match, played on patch 6.83c
func TestEsportsPatch683b(t *testing.T) {
	c := &testCase{
		matchId: 1405240741,
		url:     "http://replay135.valve.net/570/1405240741_220241732.dem.bz2",
		expectLastChatMessage: "Gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_ember_spirit": 0,
			"npc_dota_hero_broodmother":  5,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_chen":        2,
			"npc_dota_hero_broodmother": 0,
			"npc_dota_hero_sniper":      2,
			"npc_dota_hero_phoenix":     2,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84p0
func TestEsportsPatch684p0(t *testing.T) {
	c := &testCase{
		matchId: 1450235906,
		url:     "http://replay136.valve.net/570/1450235906_1463120933.dem.bz2",
		expectLastChatMessage: "gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_broodmother": 3,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_broodmother": 7,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84p1
func TestEsportsPatch684p1(t *testing.T) {
	c := &testCase{
		matchId: 1458895412,
		url:     "http://replay123.valve.net/570/1458895412_140022944.dem.bz2",
		expectLastChatMessage: "gg",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_faceless_void": 3,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_faceless_void": 2,
		},
	}

	testReplayCase(t, c)
}

// Esports match, played on patch 6.84c
func TestEsportsPatch684c(t *testing.T) {
	c := &testCase{
		matchId: 1483980562,
		url:     "http://replay133.valve.net/570/1483980562_1668922202.dem.bz2",
		expectLastChatMessage: "gg wp",
		expectHeroKillCount: map[string]int{
			"npc_dota_hero_dragon_knight": 5,
			"npc_dota_hero_bristleback":   1,
		},
		expectHeroDeathCount: map[string]int{
			"npc_dota_hero_earthshaker": 6,
			"npc_dota_hero_bristleback": 3,
		},
	}

	testReplayCase(t, c)
}

// Manually scrutinised match, played on patch 6.84p1
func TestPublicMatchPatch684p1(t *testing.T) {
	assert := assert.New(t)

	data, err := getReplayData(1456774107, "http://s.tsai.co/replays/1456774107.dem")
	if err != nil {
		t.Fatalf("unable to get replay: %s", err)
	}

	parser := NewParser(data)
	parser.OnSayText2 = func(n int, o *dota.CUserMsg_SayText2) {
	}

	earthshakerDeaths := 0
	spiritBreakerDeaths := 0
	parser.OnCombatLog = func(entry CombatLogEntry) {
		// t.Logf("OnCombatLog: %s: %+v", reflect.TypeOf(entry), entry)
		switch log := entry.(type) {
		case *CombatLogDeath:
			if log.Target == "npc_dota_hero_earthshaker" {
				earthshakerDeaths++
			}
			if log.Target == "npc_dota_hero_spirit_breaker" {
				spiritBreakerDeaths++
			}
		}
	}

	var now time.Duration
	var gameTime, preGameStarttime float64
	parser.OnEntityPreserved = func(pe *PacketEntity) {
		if pe.Name == "DT_DOTAGamerulesProxy" {
			gameTime = pe.Values["DT_DOTAGamerules.m_fGameTime"].(float64)
			preGameStarttime = pe.Values["DT_DOTAGamerules.m_flPreGameStartTime"].(float64)
			now = time.Duration(gameTime-preGameStarttime) * time.Second
		}
	}

	// entindex:3 order_type:1 units:349 position:<x:6953.3125 y:6920.8438 z:384 > queue:false
	unitOrderCount := 0
	unitOrderQueuedCount := 0
	specificUnitOrder := false
	parser.OnSpectatorPlayerUnitOrders = func(n int, o *dota.CDOTAUserMsg_SpectatorPlayerUnitOrders) {
		unitOrderCount++
		if *o.Queue == true {
			unitOrderQueuedCount++
		}
		if *o.Entindex == 3 && *o.OrderType == 1 && o.Units[0] == 349 && *o.Queue == false &&
			*o.Position.X == 6953.3125 && *o.Position.Y == 6920.8438 && *o.Position.Y == 384.0 {
			specificUnitOrder = true
		}
	}

	chatWheelMessagesCount := 0
	parser.OnChatWheel = func(n int, o *dota.CDOTAUserMsg_ChatWheel) {
		chatWheelMessagesCount++
	}

	parser.Parse()

	assert.Equal(8, earthshakerDeaths)
	assert.Equal(11, spiritBreakerDeaths)          // not actually right but verified in replay
	assert.Equal(55316, unitOrderCount)            // regression test
	assert.Equal(102, unitOrderQueuedCount)        // regression test
	assert.Equal(int64(2585000000000), int64(now)) // regression test
	assert.Equal(0, chatWheelMessagesCount)        // regression test
}

func testReplayCase(t *testing.T, c *testCase) {
	assert := assert.New(t)

	data, err := getReplayData(c.matchId, c.url)
	if err != nil {
		t.Fatalf("unable to get replay: %s", err)
	}

	worldMins := &Vector3{}
	worldMaxes := &Vector3{}
	lastChatMessage := ""
	heroKillCount := make(map[string]int)
	heroDeathCount := make(map[string]int)

	parser := NewParser(data)
	parser.OnSayText2 = func(n int, o *dota.CUserMsg_SayText2) {
		lastChatMessage = o.GetText()
	}

	parser.OnChatEvent = func(n int, o *dota.CDOTAUserMsg_ChatEvent) {
	}

	parser.OnCombatLog = func(entry CombatLogEntry) {
		switch log := entry.(type) {
		case *CombatLogDeath:
			if strings.HasPrefix(log.Target, "npc_dota_hero_") {
				if _, ok := heroKillCount[log.Source]; !ok {
					heroKillCount[log.Source] = 0
				}
				heroKillCount[log.Source] += 1
			}

			if _, ok := heroDeathCount[log.Target]; !ok {
				heroDeathCount[log.Target] = 0
			}
			heroDeathCount[log.Target] += 1
		}
	}

	parser.OnEntityCreated = func(ent *PacketEntity) {
		if ent.Tick == 0 && ent.Name == "DT_WORLD" {
			worldMins = ent.Values["DT_WORLD.m_WorldMins"].(*Vector3)
			worldMaxes = ent.Values["DT_WORLD.m_WorldMaxs"].(*Vector3)
		}
	}

	parser.Parse()

	// Make sure we have found the death counts for specified heroes
	if c.expectHeroDeathCount != nil {
		for hero, count := range c.expectHeroDeathCount {
			assert.Equal(count, heroDeathCount[hero], "expected hero %s to have death count %d", hero, count)
		}
	}

	// Make sure we have found the kill counts for specified heroes.
	if c.expectHeroKillCount != nil {
		for hero, count := range c.expectHeroKillCount {
			assert.Equal(count, heroKillCount[hero], "expected hero %s to have kill count %d", hero, count)
		}
	}

	// Make sure we find the DT_WORLD entity and it has the correct min and max dimensions.
	// This serves to help ensure our Float and Vector3 parsing is correct.
	assert.Equal(&Vector3{X: -8576.0, Y: -7680.0, Z: -1536.0}, worldMins)
	assert.Equal(&Vector3{X: 9216.0, Y: 8192.0, Z: 256.0}, worldMaxes)

	// Make sure we found the chat messages and have properly found the last one
	assert.Equal(c.expectLastChatMessage, lastChatMessage)
}

func getReplayData(matchId int64, url string) ([]byte, error) {
	path := fmt.Sprintf("replays/%d.dem", matchId)
	if data, err := ioutil.ReadFile(path); err == nil {
		fmt.Printf("read replay %d from %s\n", matchId, path)
		return data, nil
	}

	fmt.Printf("downloading replay %d from %s...\n", matchId, url)

	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	// Return an error if we don't get a 200
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("invalid status %d", resp.StatusCode)
	}

	var data []byte
	if url[len(url)-3:] == "bz2" {
		data, err = ioutil.ReadAll(bzip2.NewReader(resp.Body))
	} else {
		data, err = ioutil.ReadAll(resp.Body)
	}

	if err != nil {
		return nil, err
	}

	if err := ioutil.WriteFile(path, data, 0644); err != nil {
		return nil, err
	}

	fmt.Printf("downloaded replay %d from %s to %s\n", matchId, url, path)

	return data, nil
}
<|endoftext|>


<commit_before>package watchdog

import (
	"fmt"
	"github.com/appio/watchdog/process"
	"sync"
)

// The purpose of this package is to act as a registry for processes
// and proxy commands to processes and respond to events from processes, even
// if the event handling is only logging events.
//
// It is also responsible for sending log output to drain channels.
//
// All exported methods in this package are designed to be interacted with by the `agent` package.
//
// In typical operation this package would only be run once as a daemon per host.

type Watchdog struct {
	childProcesses map[string]*process.Process
	pMu            sync.Mutex
}

func New() *Watchdog {
	return &Watchdog{
		childProcesses: make(map[string]*process.Process),
	}
}

// Add a process
func (w *Watchdog) Add(p *process.Process) error {
	w.pMu.Lock()
	defer w.pMu.Unlock()

	if _, exists := w.childProcesses[p.Name]; exists {
		return fmt.Errorf("process already exists: %s", p.Name)
	}

	w.childProcesses[p.Name] = p

	return nil
}

// Remove a process
func (w *Watchdog) Remove(p *process.Process) error {
	w.pMu.Lock()
	defer w.pMu.Unlock()

	if _, exists := w.childProcesses[p.Name]; !exists {
		return fmt.Errorf("process not found: %s", p.Name)
	}

	delete(w.childProcesses, p.Name)

	return nil
}

// FindByName returns a process for the given name or nil if not found
func (w *Watchdog) FindByName(name string) *process.Process {
	w.pMu.Lock()
	defer w.pMu.Unlock()

	return w.childProcesses[name]
}
<commit_msg>Handle process output<commit_after>package watchdog

import (
	"fmt"
	"github.com/appio/watchdog/process"
	"strings"
	"sync"
)

// The purpose of this package is to act as a registry for processes
// and proxy commands to processes and respond to events from processes, even
// if the event handling is only logging events.
//
// It is also responsible for sending log output to drain channels.
//
// All exported methods in this package are designed to be interacted with by the `agent` package.
//
// In typical operation this package would only be run once as a daemon per host.

type Watchdog struct {
	childProcesses map[string]*process.Process
	managed        map[string]chan bool
	pMu            sync.Mutex
	manage         chan int
}

func New() *Watchdog {
	return &Watchdog{
		childProcesses: make(map[string]*process.Process),
		managed:        make(map[string]chan bool, 1),
		manage:         make(chan int),
	}
}

// Add a process
func (w *Watchdog) Add(p *process.Process) error {
	w.pMu.Lock()
	defer w.pMu.Unlock()

	if _, exists := w.childProcesses[p.Name]; exists {
		return fmt.Errorf("process already exists: %s", p.Name)
	}

	w.childProcesses[p.Name] = p
	w.manageProcess(p)

	return nil
}

// Remove a process
func (w *Watchdog) Remove(p *process.Process) error {
	w.pMu.Lock()
	defer w.pMu.Unlock()

	if _, exists := w.childProcesses[p.Name]; !exists {
		return fmt.Errorf("process not found: %s", p.Name)
	}

	w.managed[p.Name] <- true

	delete(w.childProcesses, p.Name)
	delete(w.managed, p.Name)

	return nil
}

// FindByName returns a process for the given name or nil if not found
func (w *Watchdog) FindByName(name string) *process.Process {
	w.pMu.Lock()
	defer w.pMu.Unlock()

	return w.childProcesses[name]
}

func (w *Watchdog) manageProcess(p *process.Process) error {
	w.managed[p.Name] = make(chan bool)

	go func() {
		for {
			select {
			case <-w.managed[p.Name]:
				return

			case out := <-p.OutputChan():
				fmt.Printf("[%s] > %s\n", p.Name, strings.TrimRight(string(out), "\n"))
			}
		}
	}()

	return nil
}
<|endoftext|>


<commit_before>// +build freebsd

package memory

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"strings"
)

// Get memory statistics
func Get() (*Stats, error) {
	cmd := exec.Command("top", "-b", "-n", "1")
	out, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}
	memory, err := collectMemoryStats(out)
	if err != nil {
		return nil, err
	}
	if err := cmd.Wait(); err != nil {
		return nil, err
	}
	return memory, nil
}

// Stats represents memory statistics for freebsd
type Stats struct {
	Total, Used, Buffers, Free, Active, Inactive, Wired,
	SwapTotal, SwapFree uint64
}

func collectMemoryStats(out io.Reader) (*Stats, error) {
	scanner := bufio.NewScanner(out)

	var memory Stats
	memStats := map[string]*uint64{
		"MemBuf":    &memory.Buffers,
		"MemFree":   &memory.Free,
		"MemActive": &memory.Active,
		"MemInact":  &memory.Inactive,
		"MemWired":  &memory.Wired,
		"SwapTotal": &memory.SwapTotal,
		"SwapFree":  &memory.SwapFree,
	}

	var cnt int
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "Mem:") && !strings.HasPrefix(line, "Swap:") {
			continue
		}
		cnt += 1
		i := strings.IndexRune(line, ':')
		prefix := line[:i]
		stats := strings.Split(line[i+1:], ",")
		for _, stat := range stats {
			cs := strings.Fields(stat)
			if len(cs) != 2 {
				continue
			}
			if ptr := memStats[prefix+cs[1]]; ptr != nil {
				if val, err := parseValue(cs[0]); err == nil {
					*ptr = val
				}
			}
		}
		if cnt == 2 {
			break
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("scan error for top -b -n 1: %s", err)
	}

	return &memory, nil
}

func parseValue(valStr string) (uint64, error) {
	if len(valStr) < 1 {
		return 0, errors.New("empty value")
	}
	var unit uint64
	switch valStr[len(valStr)-1] {
	case 'T':
		unit = 1024 * 1024 * 1024 * 1024
	case 'G':
		unit = 1024 * 1024 * 1024
	case 'M':
		unit = 1024 * 1024
	case 'K':
		unit = 1024
	default:
		unit = 1
	}
	val, err := strconv.ParseUint(valStr[:len(valStr)-1], 10, 64)
	return val * unit, err
}
<commit_msg>add total memory for freebsd<commit_after>// +build freebsd

package memory

import (
	"bufio"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
)

// Get memory statistics
func Get() (*Stats, error) {
	cmd := exec.Command("top", "-b", "-n", "1")
	out, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}
	memory, err := collectMemoryStats(out)
	if err != nil {
		return nil, err
	}
	if err := cmd.Wait(); err != nil {
		return nil, err
	}

	ret, err := syscall.Sysctl("hw.physmem")
	if err != nil {
		return nil, fmt.Errorf("failed in sysctl hw.physmem: %s", err)
	}
	memory.Total = binary.LittleEndian.Uint64([]byte(ret + "\x00"))

	return memory, nil
}

// Stats represents memory statistics for freebsd
type Stats struct {
	Total, Used, Buffers, Free, Active, Inactive, Wired,
	SwapTotal, SwapFree uint64
}

func collectMemoryStats(out io.Reader) (*Stats, error) {
	scanner := bufio.NewScanner(out)

	var memory Stats
	memStats := map[string]*uint64{
		"MemBuf":    &memory.Buffers,
		"MemFree":   &memory.Free,
		"MemActive": &memory.Active,
		"MemInact":  &memory.Inactive,
		"MemWired":  &memory.Wired,
		"SwapTotal": &memory.SwapTotal,
		"SwapFree":  &memory.SwapFree,
	}

	var cnt int
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "Mem:") && !strings.HasPrefix(line, "Swap:") {
			continue
		}
		cnt += 1
		i := strings.IndexRune(line, ':')
		prefix := line[:i]
		stats := strings.Split(line[i+1:], ",")
		for _, stat := range stats {
			cs := strings.Fields(stat)
			if len(cs) != 2 {
				continue
			}
			if ptr := memStats[prefix+cs[1]]; ptr != nil {
				if val, err := parseValue(cs[0]); err == nil {
					*ptr = val
				}
			}
		}
		if cnt == 2 {
			break
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("scan error for top -b -n 1: %s", err)
	}

	return &memory, nil
}

func parseValue(valStr string) (uint64, error) {
	if len(valStr) < 1 {
		return 0, errors.New("empty value")
	}
	var unit uint64
	switch valStr[len(valStr)-1] {
	case 'T':
		unit = 1024 * 1024 * 1024 * 1024
	case 'G':
		unit = 1024 * 1024 * 1024
	case 'M':
		unit = 1024 * 1024
	case 'K':
		unit = 1024
	default:
		unit = 1
	}
	val, err := strconv.ParseUint(valStr[:len(valStr)-1], 10, 64)
	return val * unit, err
}
<|endoftext|>


<commit_before>package curator_test

import (
	. "github.com/talbright/go-curator"
	"github.com/talbright/go-zookeeper/zk"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Znode", func() {
	var node *Znode
	BeforeEach(func() {
		node = NewZnode("/usr/local/var/base")
	})
	Context("NewZnode", func() {
		It("creates a node using the basename as the name", func() {
			Expect(node.Name).Should(Equal("base"))
			Expect(node.Path).Should(Equal("/usr/local/var/base"))
		})
	})
	Context("Basename", func() {
		It("should extract the basename of the node", func() {
			Expect(node.Basename()).Should(Equal("base"))
		})
	})
	Context("Parent", func() {
		It("should extract the parent name of the node", func() {
			Expect(node.Parent()).Should(Equal("/usr/local/var"))
		})
	})
	Context("DeepCopy", func() {
		It("should create a deep copy of the znode with nil data and stat", func() {
			znodeCopy := node.DeepCopy()
			Expect(znodeCopy.Data).Should(BeNil())
			Expect(znodeCopy.Stat).Should(BeNil())
		})
		It("should create a deep copy of the znode with data and stat", func() {
			node.Stat = &zk.Stat{Version: 123}
			node.Data = []byte("abc")
			znodeCopy := node.DeepCopy()
			Expect(znodeCopy.Stat == node.Stat).ShouldNot(BeTrue())
			node.Stat.Version = 99
			Expect(znodeCopy.Stat.Version).ShouldNot(Equal(node.Stat.Version))
			node.Data[0] = 'z'
			Expect(string(znodeCopy.Data[:])).ShouldNot(Equal(string(node.Data[:])))
		})
	})
})
<commit_msg>Improvements to znode test<commit_after>package curator_test

import (
	. "github.com/talbright/go-curator"
	"github.com/talbright/go-zookeeper/zk"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Znode", func() {
	var node *Znode
	BeforeEach(func() {
		node = NewZnode("/usr/local/var/base")
	})
	Context("NewZnode", func() {
		It("creates a node using the basename as the name", func() {
			Expect(node.Name).Should(Equal("base"))
			Expect(node.Path).Should(Equal("/usr/local/var/base"))
		})
	})
	Context("Basename", func() {
		It("should extract the basename of the node", func() {
			Expect(node.Basename()).Should(Equal("base"))
		})
	})
	Context("Parent", func() {
		It("should extract the parent name of the node", func() {
			Expect(node.Parent()).Should(Equal("/usr/local/var"))
		})
	})
	Context("DeepCopy", func() {
		It("should create a deep copy of the znode with nil data and stat", func() {
			znodeCopy := node.DeepCopy()
			Expect(znodeCopy.Data).Should(BeNil())
			Expect(znodeCopy.Stat).Should(BeNil())
		})
		It("should create a deep copy of the znode with data and stat", func() {
			node.Stat = &zk.Stat{Version: 123}
			node.Data = []byte("abc")
			znodeCopy := node.DeepCopy()
			Expect(znodeCopy.Name).Should(Equal(node.Name))
			Expect(znodeCopy.Path).Should(Equal(node.Path))
			Expect(znodeCopy.Stat).Should(Equal(node.Stat))
			node.Stat.Version = 99
			Expect(znodeCopy.Stat.Version).ShouldNot(Equal(node.Stat.Version))
			Expect(string(znodeCopy.Data[:])).Should(Equal(string(node.Data[:])))
			node.Data[0] = 'z'
			Expect(string(znodeCopy.Data[:])).ShouldNot(Equal(string(node.Data[:])))
		})
	})
})
<|endoftext|>


<commit_before>package collectors

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/amonapp/amonagent/logging"
	"github.com/amonapp/amonagent/plugins"
	"github.com/amonapp/amonagent/settings"
)

// CollectorLogger - XXX
var CollectorLogger = logging.GetLogger("amonagent.collector")

func (p SystemDataStruct) String() string {
	s, _ := json.Marshal(p)
	return string(s)
}

func (p AllMetricsStruct) String() string {
	s, _ := json.Marshal(p)
	return string(s)

}

func (p HostDataStruct) String() string {
	s, _ := json.Marshal(p)
	return string(s)
}

// AllMetricsStruct -XXX
type AllMetricsStruct struct {
	System    SystemDataStruct `json:"system"`
	Processes ProcessesList    `json:"processes"`
	Host      HostDataStruct   `json:"host"`
	Plugins   interface{}      `json:"plugins"`
	Checks    interface{}      `json:"checks"`
}

// HostDataStruct -XXX
type HostDataStruct struct {
	Host       string       `json:"host"`
	MachineID  string       `json:"machineid"`
	ServerKey  string       `json:"server_key"`
	Distro     DistroStruct `json:"distro"`
	IPAddress  string       `json:"ip_address"`
	InstanceID string       `json:"instance_id"`
}

// SystemDataStruct - collect all system metrics
type SystemDataStruct struct {
	CPU     CPUUsageStruct   `json:"cpu"`
	Network NetworkUsageList `json:"network"`
	Disk    DiskUsageList    `json:"disk"`
	Load    LoadStruct       `json:"loadavg"`
	Uptime  string           `json:"uptime"`
	Memory  MemoryStruct     `json:"memory"`
}

// CollectPluginsData - XXX
func CollectPluginsData() (interface{}, interface{}) {
	PluginResults := make(map[string]interface{})
	var CheckResults interface{}
	var wg sync.WaitGroup
	EnabledPlugins, _ := plugins.GetAllEnabledPlugins()

	resultChan := make(chan interface{}, len(EnabledPlugins))

	for _, p := range EnabledPlugins {
		wg.Add(1)
		creator, _ := plugins.Plugins[p.Name]
		plugin := creator()

		go func(p plugins.PluginConfig) {
			PluginResult, err := plugin.Collect(p.Path)
			if err != nil {
				CollectorLogger.Errorf("Can't get stats for plugin: %s", err)

			}

			resultChan <- PluginResult
			defer wg.Done()
		}(p)

		// if p.Name == "checks" {
		// 	CheckResults = resultChan
		// } else {
		// 	PluginResults[p.Name] = resultChan
		// }

	}

	wg.Wait()
	close(resultChan)

	for i := range resultChan {
		fmt.Println(i)
		fmt.Println("---------------------------------------------------")
		// result = append(result, i)
	}

	return PluginResults, CheckResults
}

// CollectHostData - XXX
func CollectHostData() HostDataStruct {

	host := Host()
	// Load settings
	settings := settings.Settings()

	var machineID string
	var InstanceID string
	var ip string
	var distro DistroStruct

	machineID = GetOrCreateMachineID()
	InstanceID = CloudID()
	ip = IPAddress()
	distro = Distro()

	hoststruct := HostDataStruct{
		Host:       host,
		MachineID:  machineID,
		Distro:     distro,
		IPAddress:  ip,
		ServerKey:  settings.ServerKey,
		InstanceID: InstanceID,
	}

	return hoststruct
}

// CollectSystemData - XXX
func CollectSystemData() SystemDataStruct {
	var networkUsage NetworkUsageList
	var cpuUsage CPUUsageStruct
	var diskUsage DiskUsageList
	var memoryUsage MemoryStruct
	var UptimeString string
	var Load LoadStruct

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		networkUsage, _ = NetworkUsage()
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		cpuUsage = CPUUsage()
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		diskUsage, _ = DiskUsage()
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		memoryUsage = MemoryUsage()
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		UptimeString = Uptime()
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		Load = LoadAverage()
	}()

	wg.Wait()

	SystemData := SystemDataStruct{
		CPU:     cpuUsage,
		Network: networkUsage,
		Disk:    diskUsage,
		Load:    Load,
		Uptime:  UptimeString,
		Memory:  memoryUsage,
	}

	return SystemData

}

// CollectProcessData - XXX
func CollectProcessData() ProcessesList {
	var ProcessesUsage ProcessesList
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		ProcessesUsage, _ = Processes()
	}()

	wg.Wait()

	return ProcessesUsage
}

// CollectAllData - XXX
func CollectAllData() AllMetricsStruct {

	ProcessesData := CollectProcessData()
	SystemData := CollectSystemData()
	Plugins, Checks := CollectPluginsData()
	HostData := CollectHostData()

	allMetrics := AllMetricsStruct{
		System:    SystemData,
		Processes: ProcessesData,
		Host:      HostData,
		Plugins:   Plugins,
		Checks:    Checks,
	}

	return allMetrics
}
<commit_msg>Fix plugin collector data race<commit_after>package collectors

import (
	"encoding/json"
	"sync"

	"github.com/amonapp/amonagent/logging"
	"github.com/amonapp/amonagent/plugins"
	"github.com/amonapp/amonagent/settings"
)

// CollectorLogger - XXX
var CollectorLogger = logging.GetLogger("amonagent.collector")

func (p SystemDataStruct) String() string {
	s, _ := json.Marshal(p)
	return string(s)
}

func (p AllMetricsStruct) String() string {
	s, _ := json.Marshal(p)
	return string(s)

}

func (p HostDataStruct) String() string {
	s, _ := json.Marshal(p)
	return string(s)
}

// AllMetricsStruct -XXX
type AllMetricsStruct struct {
	System    SystemDataStruct `json:"system"`
	Processes ProcessesList    `json:"processes"`
	Host      HostDataStruct   `json:"host"`
	Plugins   interface{}      `json:"plugins"`
	Checks    interface{}      `json:"checks"`
}

// HostDataStruct -XXX
type HostDataStruct struct {
	Host       string       `json:"host"`
	MachineID  string       `json:"machineid"`
	ServerKey  string       `json:"server_key"`
	Distro     DistroStruct `json:"distro"`
	IPAddress  string       `json:"ip_address"`
	InstanceID string       `json:"instance_id"`
}

// SystemDataStruct - collect all system metrics
type SystemDataStruct struct {
	CPU     CPUUsageStruct   `json:"cpu"`
	Network NetworkUsageList `json:"network"`
	Disk    DiskUsageList    `json:"disk"`
	Load    LoadStruct       `json:"loadavg"`
	Uptime  string           `json:"uptime"`
	Memory  MemoryStruct     `json:"memory"`
}

// PluginResultStruct - a channel struct that holds plugin results
type PluginResultStruct struct {
	Name   string
	Result interface{}
}

// CollectPluginsData - XXX
func CollectPluginsData() (interface{}, interface{}) {
	PluginResults := make(map[string]interface{})
	var CheckResults interface{}
	var wg sync.WaitGroup
	EnabledPlugins, _ := plugins.GetAllEnabledPlugins()

	resultChan := make(chan PluginResultStruct, len(EnabledPlugins))

	for _, p := range EnabledPlugins {
		wg.Add(1)
		creator, _ := plugins.Plugins[p.Name]
		plugin := creator()

		go func(p plugins.PluginConfig) {
			PluginResult, err := plugin.Collect(p.Path)
			if err != nil {
				CollectorLogger.Errorf("Can't get stats for plugin: %s", err)
			}

			r := PluginResultStruct{Name: p.Name, Result: PluginResult}

			resultChan <- r
			defer wg.Done()
		}(p)

	}

	wg.Wait()
	close(resultChan)

	for result := range resultChan {
		if result.Name == "checks" {
			CheckResults = result.Result
		} else {
			PluginResults[result.Name] = result.Result
		}

	}

	return PluginResults, CheckResults
}

// CollectHostData - XXX
func CollectHostData() HostDataStruct {

	host := Host()
	// Load settings
	settings := settings.Settings()

	var machineID string
	var InstanceID string
	var ip string
	var distro DistroStruct

	machineID = GetOrCreateMachineID()
	InstanceID = CloudID()
	ip = IPAddress()
	distro = Distro()

	hoststruct := HostDataStruct{
		Host:       host,
		MachineID:  machineID,
		Distro:     distro,
		IPAddress:  ip,
		ServerKey:  settings.ServerKey,
		InstanceID: InstanceID,
	}

	return hoststruct
}

// CollectSystemData - XXX
func CollectSystemData() SystemDataStruct {
	var networkUsage NetworkUsageList
	var cpuUsage CPUUsageStruct
	var diskUsage DiskUsageList
	var memoryUsage MemoryStruct
	var UptimeString string
	var Load LoadStruct

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		networkUsage, _ = NetworkUsage()
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		cpuUsage = CPUUsage()
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		diskUsage, _ = DiskUsage()
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		memoryUsage = MemoryUsage()
	}()
	wg.Add(1)
	go func() {
		defer wg.Done()
		UptimeString = Uptime()
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		Load = LoadAverage()
	}()

	wg.Wait()

	SystemData := SystemDataStruct{
		CPU:     cpuUsage,
		Network: networkUsage,
		Disk:    diskUsage,
		Load:    Load,
		Uptime:  UptimeString,
		Memory:  memoryUsage,
	}

	return SystemData

}

// CollectProcessData - XXX
func CollectProcessData() ProcessesList {
	var ProcessesUsage ProcessesList
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		ProcessesUsage, _ = Processes()
	}()

	wg.Wait()

	return ProcessesUsage
}

// CollectAllData - XXX
func CollectAllData() AllMetricsStruct {

	ProcessesData := CollectProcessData()
	SystemData := CollectSystemData()
	Plugins, Checks := CollectPluginsData()
	HostData := CollectHostData()

	allMetrics := AllMetricsStruct{
		System:    SystemData,
		Processes: ProcessesData,
		Host:      HostData,
		Plugins:   Plugins,
		Checks:    Checks,
	}

	return allMetrics
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors 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 collectors

import (
	"encoding/json"
	"regexp"
	"strconv"

	"github.com/golang/glog"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/net/context"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"
)

var (
	invalidLabelCharRE         = regexp.MustCompile(`[^a-zA-Z0-9_]`)
	descPodLabelsName          = "kube_pod_labels"
	descPodLabelsHelp          = "Kubernetes labels converted to Prometheus labels."
	descPodLabelsDefaultLabels = []string{"namespace", "pod"}
	containerWaitingReasons    = []string{"ContainerCreating", "ErrImagePull"}

	descPodInfo = prometheus.NewDesc(
		"kube_pod_info",
		"Information about pod.",
		[]string{"namespace", "pod", "host_ip", "pod_ip", "node", "created_by_kind", "created_by_name"}, nil,
	)

	descPodStartTime = prometheus.NewDesc(
		"kube_pod_start_time",
		"Start time in unix timestamp for a pod.",
		[]string{"namespace", "pod"}, nil,
	)

	descPodOwner = prometheus.NewDesc(
		"kube_pod_owner",
		"Information about the Pod's owner.",
		[]string{"namespace", "pod", "owner_kind", "owner_name", "owner_is_controller"}, nil,
	)

	descPodLabels = prometheus.NewDesc(
		descPodLabelsName,
		descPodLabelsHelp,
		descPodLabelsDefaultLabels, nil,
	)

	descPodCreated = prometheus.NewDesc(
		"kube_pod_created",
		"Unix creation timestamp",
		[]string{"namespace", "pod"}, nil,
	)

	descPodStatusPhase = prometheus.NewDesc(
		"kube_pod_status_phase",
		"The pods current phase.",
		[]string{"namespace", "pod", "phase"}, nil,
	)

	descPodStatusReady = prometheus.NewDesc(
		"kube_pod_status_ready",
		"Describes whether the pod is ready to serve requests.",
		[]string{"namespace", "pod", "condition"}, nil,
	)

	descPodStatusScheduled = prometheus.NewDesc(
		"kube_pod_status_scheduled",
		"Describes the status of the scheduling process for the pod.",
		[]string{"namespace", "pod", "condition"}, nil,
	)

	descPodContainerInfo = prometheus.NewDesc(
		"kube_pod_container_info",
		"Information about a container in a pod.",
		[]string{"namespace", "pod", "container", "image", "image_id", "container_id"}, nil,
	)

	descPodContainerStatusWaiting = prometheus.NewDesc(
		"kube_pod_container_status_waiting",
		"Describes whether the container is currently in waiting state.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerStatusWaitingReason = prometheus.NewDesc(
		"kube_pod_container_status_waiting_reason",
		"Describes the reason the container is currently in waiting state.",
		[]string{"namespace", "pod", "container", "reason"}, nil,
	)

	descPodContainerStatusRunning = prometheus.NewDesc(
		"kube_pod_container_status_running",
		"Describes whether the container is currently in running state.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerStatusTerminated = prometheus.NewDesc(
		"kube_pod_container_status_terminated",
		"Describes whether the container is currently in terminated state.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerStatusReady = prometheus.NewDesc(
		"kube_pod_container_status_ready",
		"Describes whether the containers readiness check succeeded.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerStatusRestarts = prometheus.NewDesc(
		"kube_pod_container_status_restarts",
		"The number of container restarts per container.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerResourceRequestsCpuCores = prometheus.NewDesc(
		"kube_pod_container_resource_requests_cpu_cores",
		"The number of requested cpu cores by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceRequestsMemoryBytes = prometheus.NewDesc(
		"kube_pod_container_resource_requests_memory_bytes",
		"The number of requested memory bytes  by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceLimitsCpuCores = prometheus.NewDesc(
		"kube_pod_container_resource_limits_cpu_cores",
		"The limit on cpu cores to be used by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceLimitsMemoryBytes = prometheus.NewDesc(
		"kube_pod_container_resource_limits_memory_bytes",
		"The limit on memory to be used by a container in bytes.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceRequestsNvidiaGPUDevices = prometheus.NewDesc(
		"kube_pod_container_resource_requests_nvidia_gpu_devices",
		"The number of requested gpu devices by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceLimitsNvidiaGPUDevices = prometheus.NewDesc(
		"kube_pod_container_resource_limits_nvidia_gpu_devices",
		"The limit on gpu devices to be used by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)
)

type PodLister func() ([]v1.Pod, error)

func (l PodLister) List() ([]v1.Pod, error) {
	return l()
}

func RegisterPodCollector(registry prometheus.Registerer, kubeClient kubernetes.Interface, namespace string) {
	client := kubeClient.CoreV1().RESTClient()
	plw := cache.NewListWatchFromClient(client, "pods", namespace, nil)
	pinf := cache.NewSharedInformer(plw, &v1.Pod{}, resyncPeriod)

	podLister := PodLister(func() (pods []v1.Pod, err error) {
		for _, m := range pinf.GetStore().List() {
			pods = append(pods, *m.(*v1.Pod))
		}
		return pods, nil
	})

	registry.MustRegister(&podCollector{store: podLister})
	go pinf.Run(context.Background().Done())
}

type podStore interface {
	List() (pods []v1.Pod, err error)
}

// podCollector collects metrics about all pods in the cluster.
type podCollector struct {
	store podStore
}

// Describe implements the prometheus.Collector interface.
func (pc *podCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- descPodInfo
	ch <- descPodStartTime
	ch <- descPodOwner
	ch <- descPodLabels
	ch <- descPodCreated
	ch <- descPodStatusPhase
	ch <- descPodStatusReady
	ch <- descPodStatusScheduled
	ch <- descPodContainerInfo
	ch <- descPodContainerStatusWaiting
	ch <- descPodContainerStatusWaitingReason
	ch <- descPodContainerStatusRunning
	ch <- descPodContainerStatusTerminated
	ch <- descPodContainerStatusReady
	ch <- descPodContainerStatusRestarts
	ch <- descPodContainerResourceRequestsCpuCores
	ch <- descPodContainerResourceRequestsMemoryBytes
	ch <- descPodContainerResourceLimitsCpuCores
	ch <- descPodContainerResourceLimitsMemoryBytes
	ch <- descPodContainerResourceRequestsNvidiaGPUDevices
	ch <- descPodContainerResourceLimitsNvidiaGPUDevices
}

func extractCreatedBy(annotation map[string]string) *api.ObjectReference {
	value, ok := annotation[api.CreatedByAnnotation]
	if ok {
		var r api.SerializedReference
		err := json.Unmarshal([]byte(value), &r)
		if err == nil {
			return &r.Reference
		}
	}
	return nil
}

// Collect implements the prometheus.Collector interface.
func (pc *podCollector) Collect(ch chan<- prometheus.Metric) {
	pods, err := pc.store.List()
	if err != nil {
		glog.Errorf("listing pods failed: %s", err)
		return
	}
	for _, p := range pods {
		pc.collectPod(ch, p)
	}

	glog.Infof("collected %d pods", len(pods))
}

func kubeLabelsToPrometheusLabels(labels map[string]string) ([]string, []string) {
	labelKeys := make([]string, len(labels))
	labelValues := make([]string, len(labels))
	i := 0
	for k, v := range labels {
		labelKeys[i] = "label_" + sanitizeLabelName(k)
		labelValues[i] = v
		i++
	}
	return labelKeys, labelValues
}

func sanitizeLabelName(s string) string {
	return invalidLabelCharRE.ReplaceAllString(s, "_")
}

func podLabelsDesc(labelKeys []string) *prometheus.Desc {
	return prometheus.NewDesc(
		descPodLabelsName,
		descPodLabelsHelp,
		append(descPodLabelsDefaultLabels, labelKeys...),
		nil,
	)
}

func (pc *podCollector) collectPod(ch chan<- prometheus.Metric, p v1.Pod) {
	nodeName := p.Spec.NodeName
	addConstMetric := func(desc *prometheus.Desc, t prometheus.ValueType, v float64, lv ...string) {
		lv = append([]string{p.Namespace, p.Name}, lv...)
		ch <- prometheus.MustNewConstMetric(desc, t, v, lv...)
	}
	addGauge := func(desc *prometheus.Desc, v float64, lv ...string) {
		addConstMetric(desc, prometheus.GaugeValue, v, lv...)
	}
	addCounter := func(desc *prometheus.Desc, v float64, lv ...string) {
		addConstMetric(desc, prometheus.CounterValue, v, lv...)
	}

	createdBy := extractCreatedBy(p.Annotations)
	createdByKind := "<none>"
	createdByName := "<none>"
	if createdBy != nil {
		if createdBy.Kind != "" {
			createdByKind = createdBy.Kind
		}
		if createdBy.Name != "" {
			createdByName = createdBy.Name
		}
	}

	if p.Status.StartTime != nil {
		addGauge(descPodStartTime, float64((*(p.Status.StartTime)).Unix()))
	}

	addGauge(descPodInfo, 1, p.Status.HostIP, p.Status.PodIP, nodeName, createdByKind, createdByName)

	owners := p.GetOwnerReferences()
	if len(owners) == 0 {
		addGauge(descPodOwner, 1, "<none>", "<none>", "<none>")
	} else {
		for _, owner := range owners {
			if owner.Controller != nil {
				addGauge(descPodOwner, 1, owner.Kind, owner.Name, strconv.FormatBool(*owner.Controller))
			} else {
				addGauge(descPodOwner, 1, owner.Kind, owner.Name, "false")
			}
		}
	}

	labelKeys, labelValues := kubeLabelsToPrometheusLabels(p.Labels)
	addGauge(podLabelsDesc(labelKeys), 1, labelValues...)

	if p := p.Status.Phase; p != "" {
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodPending), string(v1.PodPending))
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodRunning), string(v1.PodRunning))
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodSucceeded), string(v1.PodSucceeded))
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodFailed), string(v1.PodFailed))
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodUnknown), string(v1.PodUnknown))
	}

	if !p.CreationTimestamp.IsZero() {
		addGauge(descPodCreated, float64(p.CreationTimestamp.Unix()))
	}

	for _, c := range p.Status.Conditions {
		switch c.Type {
		case v1.PodReady:
			addConditionMetrics(ch, descPodStatusReady, c.Status, p.Namespace, p.Name)
		case v1.PodScheduled:
			addConditionMetrics(ch, descPodStatusScheduled, c.Status, p.Namespace, p.Name)
		}
	}

	waitingReason := func(cs v1.ContainerStatus) string {
		if cs.State.Waiting == nil {
			return ""
		}
		return cs.State.Waiting.Reason
	}

	for _, cs := range p.Status.ContainerStatuses {
		addGauge(descPodContainerInfo, 1,
			cs.Name, cs.Image, cs.ImageID, cs.ContainerID,
		)
		addGauge(descPodContainerStatusWaiting, boolFloat64(cs.State.Waiting != nil), cs.Name)
		for _, reason := range containerWaitingReasons {
			addGauge(descPodContainerStatusWaitingReason, boolFloat64(cs.State.Waiting != nil && reason == waitingReason(cs)), cs.Name, reason)
		}
		addGauge(descPodContainerStatusRunning, boolFloat64(cs.State.Running != nil), cs.Name)
		addGauge(descPodContainerStatusTerminated, boolFloat64(cs.State.Terminated != nil), cs.Name)
		addGauge(descPodContainerStatusReady, boolFloat64(cs.Ready), cs.Name)
		addCounter(descPodContainerStatusRestarts, float64(cs.RestartCount), cs.Name)
	}

	for _, c := range p.Spec.Containers {
		req := c.Resources.Requests
		lim := c.Resources.Limits

		if cpu, ok := req[v1.ResourceCPU]; ok {
			addGauge(descPodContainerResourceRequestsCpuCores, float64(cpu.MilliValue())/1000,
				c.Name, nodeName)
		}
		if mem, ok := req[v1.ResourceMemory]; ok {
			addGauge(descPodContainerResourceRequestsMemoryBytes, float64(mem.Value()),
				c.Name, nodeName)
		}

		if gpu, ok := req[v1.ResourceNvidiaGPU]; ok {
			addGauge(descPodContainerResourceRequestsNvidiaGPUDevices, float64(gpu.Value()), c.Name, nodeName)
		}

		if cpu, ok := lim[v1.ResourceCPU]; ok {
			addGauge(descPodContainerResourceLimitsCpuCores, float64(cpu.MilliValue())/1000,
				c.Name, nodeName)
		}

		if mem, ok := lim[v1.ResourceMemory]; ok {
			addGauge(descPodContainerResourceLimitsMemoryBytes, float64(mem.Value()),
				c.Name, nodeName)
		}

		if gpu, ok := lim[v1.ResourceNvidiaGPU]; ok {
			addGauge(descPodContainerResourceLimitsNvidiaGPUDevices, float64(gpu.Value()), c.Name, nodeName)
		}
	}
}
<commit_msg>improved code a bit<commit_after>/*
Copyright 2016 The Kubernetes Authors 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 collectors

import (
	"encoding/json"
	"regexp"
	"strconv"

	"github.com/golang/glog"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/net/context"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"
)

var (
	invalidLabelCharRE         = regexp.MustCompile(`[^a-zA-Z0-9_]`)
	descPodLabelsName          = "kube_pod_labels"
	descPodLabelsHelp          = "Kubernetes labels converted to Prometheus labels."
	descPodLabelsDefaultLabels = []string{"namespace", "pod"}
	containerWaitingReasons    = []string{"ContainerCreating", "ErrImagePull"}

	descPodInfo = prometheus.NewDesc(
		"kube_pod_info",
		"Information about pod.",
		[]string{"namespace", "pod", "host_ip", "pod_ip", "node", "created_by_kind", "created_by_name"}, nil,
	)

	descPodStartTime = prometheus.NewDesc(
		"kube_pod_start_time",
		"Start time in unix timestamp for a pod.",
		[]string{"namespace", "pod"}, nil,
	)

	descPodOwner = prometheus.NewDesc(
		"kube_pod_owner",
		"Information about the Pod's owner.",
		[]string{"namespace", "pod", "owner_kind", "owner_name", "owner_is_controller"}, nil,
	)

	descPodLabels = prometheus.NewDesc(
		descPodLabelsName,
		descPodLabelsHelp,
		descPodLabelsDefaultLabels, nil,
	)

	descPodCreated = prometheus.NewDesc(
		"kube_pod_created",
		"Unix creation timestamp",
		[]string{"namespace", "pod"}, nil,
	)

	descPodStatusPhase = prometheus.NewDesc(
		"kube_pod_status_phase",
		"The pods current phase.",
		[]string{"namespace", "pod", "phase"}, nil,
	)

	descPodStatusReady = prometheus.NewDesc(
		"kube_pod_status_ready",
		"Describes whether the pod is ready to serve requests.",
		[]string{"namespace", "pod", "condition"}, nil,
	)

	descPodStatusScheduled = prometheus.NewDesc(
		"kube_pod_status_scheduled",
		"Describes the status of the scheduling process for the pod.",
		[]string{"namespace", "pod", "condition"}, nil,
	)

	descPodContainerInfo = prometheus.NewDesc(
		"kube_pod_container_info",
		"Information about a container in a pod.",
		[]string{"namespace", "pod", "container", "image", "image_id", "container_id"}, nil,
	)

	descPodContainerStatusWaiting = prometheus.NewDesc(
		"kube_pod_container_status_waiting",
		"Describes whether the container is currently in waiting state.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerStatusWaitingReason = prometheus.NewDesc(
		"kube_pod_container_status_waiting_reason",
		"Describes the reason the container is currently in waiting state.",
		[]string{"namespace", "pod", "container", "reason"}, nil,
	)

	descPodContainerStatusRunning = prometheus.NewDesc(
		"kube_pod_container_status_running",
		"Describes whether the container is currently in running state.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerStatusTerminated = prometheus.NewDesc(
		"kube_pod_container_status_terminated",
		"Describes whether the container is currently in terminated state.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerStatusReady = prometheus.NewDesc(
		"kube_pod_container_status_ready",
		"Describes whether the containers readiness check succeeded.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerStatusRestarts = prometheus.NewDesc(
		"kube_pod_container_status_restarts",
		"The number of container restarts per container.",
		[]string{"namespace", "pod", "container"}, nil,
	)

	descPodContainerResourceRequestsCpuCores = prometheus.NewDesc(
		"kube_pod_container_resource_requests_cpu_cores",
		"The number of requested cpu cores by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceRequestsMemoryBytes = prometheus.NewDesc(
		"kube_pod_container_resource_requests_memory_bytes",
		"The number of requested memory bytes  by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceLimitsCpuCores = prometheus.NewDesc(
		"kube_pod_container_resource_limits_cpu_cores",
		"The limit on cpu cores to be used by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceLimitsMemoryBytes = prometheus.NewDesc(
		"kube_pod_container_resource_limits_memory_bytes",
		"The limit on memory to be used by a container in bytes.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceRequestsNvidiaGPUDevices = prometheus.NewDesc(
		"kube_pod_container_resource_requests_nvidia_gpu_devices",
		"The number of requested gpu devices by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)

	descPodContainerResourceLimitsNvidiaGPUDevices = prometheus.NewDesc(
		"kube_pod_container_resource_limits_nvidia_gpu_devices",
		"The limit on gpu devices to be used by a container.",
		[]string{"namespace", "pod", "container", "node"}, nil,
	)
)

type PodLister func() ([]v1.Pod, error)

func (l PodLister) List() ([]v1.Pod, error) {
	return l()
}

func RegisterPodCollector(registry prometheus.Registerer, kubeClient kubernetes.Interface, namespace string) {
	client := kubeClient.CoreV1().RESTClient()
	plw := cache.NewListWatchFromClient(client, "pods", namespace, nil)
	pinf := cache.NewSharedInformer(plw, &v1.Pod{}, resyncPeriod)

	podLister := PodLister(func() (pods []v1.Pod, err error) {
		for _, m := range pinf.GetStore().List() {
			pods = append(pods, *m.(*v1.Pod))
		}
		return pods, nil
	})

	registry.MustRegister(&podCollector{store: podLister})
	go pinf.Run(context.Background().Done())
}

type podStore interface {
	List() (pods []v1.Pod, err error)
}

// podCollector collects metrics about all pods in the cluster.
type podCollector struct {
	store podStore
}

// Describe implements the prometheus.Collector interface.
func (pc *podCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- descPodInfo
	ch <- descPodStartTime
	ch <- descPodOwner
	ch <- descPodLabels
	ch <- descPodCreated
	ch <- descPodStatusPhase
	ch <- descPodStatusReady
	ch <- descPodStatusScheduled
	ch <- descPodContainerInfo
	ch <- descPodContainerStatusWaiting
	ch <- descPodContainerStatusWaitingReason
	ch <- descPodContainerStatusRunning
	ch <- descPodContainerStatusTerminated
	ch <- descPodContainerStatusReady
	ch <- descPodContainerStatusRestarts
	ch <- descPodContainerResourceRequestsCpuCores
	ch <- descPodContainerResourceRequestsMemoryBytes
	ch <- descPodContainerResourceLimitsCpuCores
	ch <- descPodContainerResourceLimitsMemoryBytes
	ch <- descPodContainerResourceRequestsNvidiaGPUDevices
	ch <- descPodContainerResourceLimitsNvidiaGPUDevices
}

func extractCreatedBy(annotation map[string]string) *api.ObjectReference {
	value, ok := annotation[api.CreatedByAnnotation]
	if ok {
		var r api.SerializedReference
		err := json.Unmarshal([]byte(value), &r)
		if err == nil {
			return &r.Reference
		}
	}
	return nil
}

// Collect implements the prometheus.Collector interface.
func (pc *podCollector) Collect(ch chan<- prometheus.Metric) {
	pods, err := pc.store.List()
	if err != nil {
		glog.Errorf("listing pods failed: %s", err)
		return
	}
	for _, p := range pods {
		pc.collectPod(ch, p)
	}

	glog.Infof("collected %d pods", len(pods))
}

func kubeLabelsToPrometheusLabels(labels map[string]string) ([]string, []string) {
	labelKeys := make([]string, len(labels))
	labelValues := make([]string, len(labels))
	i := 0
	for k, v := range labels {
		labelKeys[i] = "label_" + sanitizeLabelName(k)
		labelValues[i] = v
		i++
	}
	return labelKeys, labelValues
}

func sanitizeLabelName(s string) string {
	return invalidLabelCharRE.ReplaceAllString(s, "_")
}

func podLabelsDesc(labelKeys []string) *prometheus.Desc {
	return prometheus.NewDesc(
		descPodLabelsName,
		descPodLabelsHelp,
		append(descPodLabelsDefaultLabels, labelKeys...),
		nil,
	)
}

func (pc *podCollector) collectPod(ch chan<- prometheus.Metric, p v1.Pod) {
	nodeName := p.Spec.NodeName
	addConstMetric := func(desc *prometheus.Desc, t prometheus.ValueType, v float64, lv ...string) {
		lv = append([]string{p.Namespace, p.Name}, lv...)
		ch <- prometheus.MustNewConstMetric(desc, t, v, lv...)
	}
	addGauge := func(desc *prometheus.Desc, v float64, lv ...string) {
		addConstMetric(desc, prometheus.GaugeValue, v, lv...)
	}
	addCounter := func(desc *prometheus.Desc, v float64, lv ...string) {
		addConstMetric(desc, prometheus.CounterValue, v, lv...)
	}

	createdBy := extractCreatedBy(p.Annotations)
	createdByKind := "<none>"
	createdByName := "<none>"
	if createdBy != nil {
		if createdBy.Kind != "" {
			createdByKind = createdBy.Kind
		}
		if createdBy.Name != "" {
			createdByName = createdBy.Name
		}
	}

	if p.Status.StartTime != nil {
		addGauge(descPodStartTime, float64((*(p.Status.StartTime)).Unix()))
	}

	addGauge(descPodInfo, 1, p.Status.HostIP, p.Status.PodIP, nodeName, createdByKind, createdByName)

	owners := p.GetOwnerReferences()
	if len(owners) == 0 {
		addGauge(descPodOwner, 1, "<none>", "<none>", "<none>")
	} else {
		for _, owner := range owners {
			if owner.Controller != nil {
				addGauge(descPodOwner, 1, owner.Kind, owner.Name, strconv.FormatBool(*owner.Controller))
			} else {
				addGauge(descPodOwner, 1, owner.Kind, owner.Name, "false")
			}
		}
	}

	labelKeys, labelValues := kubeLabelsToPrometheusLabels(p.Labels)
	addGauge(podLabelsDesc(labelKeys), 1, labelValues...)

	if p := p.Status.Phase; p != "" {
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodPending), string(v1.PodPending))
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodRunning), string(v1.PodRunning))
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodSucceeded), string(v1.PodSucceeded))
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodFailed), string(v1.PodFailed))
		addGauge(descPodStatusPhase, boolFloat64(p == v1.PodUnknown), string(v1.PodUnknown))
	}

	if !p.CreationTimestamp.IsZero() {
		addGauge(descPodCreated, float64(p.CreationTimestamp.Unix()))
	}

	for _, c := range p.Status.Conditions {
		switch c.Type {
		case v1.PodReady:
			addConditionMetrics(ch, descPodStatusReady, c.Status, p.Namespace, p.Name)
		case v1.PodScheduled:
			addConditionMetrics(ch, descPodStatusScheduled, c.Status, p.Namespace, p.Name)
		}
	}

	waitingReason := func(cs v1.ContainerStatus, reason string) bool {
		if cs.State.Waiting == nil {
			return false
		}
		return cs.State.Waiting.Reason == reason
	}

	for _, cs := range p.Status.ContainerStatuses {
		addGauge(descPodContainerInfo, 1,
			cs.Name, cs.Image, cs.ImageID, cs.ContainerID,
		)
		addGauge(descPodContainerStatusWaiting, boolFloat64(cs.State.Waiting != nil), cs.Name)
		for _, reason := range containerWaitingReasons {
			addGauge(descPodContainerStatusWaitingReason, boolFloat64(waitingReason(cs, reason)), cs.Name, reason)
		}
		addGauge(descPodContainerStatusRunning, boolFloat64(cs.State.Running != nil), cs.Name)
		addGauge(descPodContainerStatusTerminated, boolFloat64(cs.State.Terminated != nil), cs.Name)
		addGauge(descPodContainerStatusReady, boolFloat64(cs.Ready), cs.Name)
		addCounter(descPodContainerStatusRestarts, float64(cs.RestartCount), cs.Name)
	}

	for _, c := range p.Spec.Containers {
		req := c.Resources.Requests
		lim := c.Resources.Limits

		if cpu, ok := req[v1.ResourceCPU]; ok {
			addGauge(descPodContainerResourceRequestsCpuCores, float64(cpu.MilliValue())/1000,
				c.Name, nodeName)
		}
		if mem, ok := req[v1.ResourceMemory]; ok {
			addGauge(descPodContainerResourceRequestsMemoryBytes, float64(mem.Value()),
				c.Name, nodeName)
		}

		if gpu, ok := req[v1.ResourceNvidiaGPU]; ok {
			addGauge(descPodContainerResourceRequestsNvidiaGPUDevices, float64(gpu.Value()), c.Name, nodeName)
		}

		if cpu, ok := lim[v1.ResourceCPU]; ok {
			addGauge(descPodContainerResourceLimitsCpuCores, float64(cpu.MilliValue())/1000,
				c.Name, nodeName)
		}

		if mem, ok := lim[v1.ResourceMemory]; ok {
			addGauge(descPodContainerResourceLimitsMemoryBytes, float64(mem.Value()),
				c.Name, nodeName)
		}

		if gpu, ok := lim[v1.ResourceNvidiaGPU]; ok {
			addGauge(descPodContainerResourceLimitsNvidiaGPUDevices, float64(gpu.Value()), c.Name, nodeName)
		}
	}
}
<|endoftext|>


<commit_before>package zego

import (
	"bytes"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
)

type Resource struct {
	//Headers     http.Header
	Response interface{}
	Raw      string
}

type Auth struct {
	Username  string
	Password  string
	Subdomain string
}

func errHandler(err error) {
	if err != nil {
		log.Print(err)
	}
}

func api(auth Auth, meth string, path string, params string) (*Resource, error) {

	trn := &http.Transport{}

	client := &http.Client{
		Transport: trn,
	}

	var URL string

	// Check if entire URL is in path
	if strings.HasPrefix(path, "http") {
		URL = path

		// Otherwise build url from auth components
	} else {
		if strings.HasPrefix(auth.Subdomain, "http") {
			URL = auth.Subdomain + "/api/v2/" + path
		} else {
			URL = "https://" + auth.Subdomain + "/api/v2/" + path
		}
	}

	req, err := http.NewRequest(meth, URL, bytes.NewBufferString(params))
	if err != nil {
		return nil, err
	}

	req.Header.Add("Content-Type", "application/json")

	req.SetBasicAuth(auth.Username, auth.Password)

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return &Resource{Response: &resp, Raw: string(data)}, nil

}
<commit_msg>Accept access token as a method of authentication.<commit_after>package zego

import (
	"bytes"
	"io/ioutil"
	"log"
	"net/http"
	"strings"
)

type Resource struct {
	//Headers     http.Header
	Response interface{}
	Raw      string
}

type Auth struct {
	Username  	string
	Password  	string
	AccessToken string
	Subdomain 	string
}

func errHandler(err error) {
	if err != nil {
		log.Print(err)
	}
}

func api(auth Auth, meth string, path string, params string) (*Resource, error) {

	trn := &http.Transport{}

	client := &http.Client{
		Transport: trn,
	}

	var URL string

	// Check if entire URL is in path
	if strings.HasPrefix(path, "http") {
		URL = path

		// Otherwise build url from auth components
	} else {
		if strings.HasPrefix(auth.Subdomain, "http") {
			URL = auth.Subdomain + "/api/v2/" + path
		} else {
			URL = "https://" + auth.Subdomain + "/api/v2/" + path
		}
	}

	req, err := http.NewRequest(meth, URL, bytes.NewBufferString(params))
	if err != nil {
		return nil, err
	}

	req.Header.Add("Content-Type", "application/json")

	if auth.AccessToken == "" {
		req.SetBasicAuth(auth.Username, auth.Password)
	} else {
		req.Header.Add("Authorization", "Bearer " + auth.AccessToken)
	}

	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	data, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return &Resource{Response: &resp, Raw: string(data)}, nil

}
<|endoftext|>


<commit_before>package game

import (
	"context"
	"github.com/Sirupsen/logrus"
	"github.com/pkg/errors"
	"github.com/verath/archipelago/lib/common"
	"github.com/verath/archipelago/lib/game/model"
	"github.com/verath/archipelago/lib/network"
	"sync"
	"time"
)

const (
	// A timeout for the maximum length of a game before it is force-closed.
	// Used so that a client leaving a game open will not keep its resources
	// allocated forever.
	GameMaxDuration time.Duration = 45 * time.Minute
)

// The game coordinator is responsible for connecting players to
// a game. Once enough players has been found so that a game can
// be created, the game coordinator creates and starts a new game
// with the given players.
//
// The lifetime of the game coordinator is not tied to a single
// game but rather the entire lifetime of the game server.
type Coordinator struct {
	logEntry       *logrus.Entry
	clientProvider ClientProvider
	// WaitGroup for games created by the Coordinator
	gamesWG sync.WaitGroup
}

func NewCoordinator(log *logrus.Logger, clientProvider ClientProvider) (*Coordinator, error) {
	logEntry := common.ModuleLogEntry(log, "game/coordinator")
	return &Coordinator{
		logEntry:       logEntry,
		clientProvider: clientProvider,
	}, nil
}

// Starts and runs the Coordinator. This method blocks until the context is cancelled or
// an error occurs, and always returns a non-nil error.
func (c *Coordinator) Run(ctx context.Context) error {
	c.logEntry.Info("Starting")
	err := c.run(ctx)
	c.logEntry.Debug("Waiting for games to stop...")
	c.gamesWG.Wait()
	c.logEntry.Info("Stopped")
	return err
}

// Waits for two player connections to be made. If successful, the methods
// returns two started clients. These clients *must* be stopped. If the
// method returns an error the clients can be assumed to be stopped.
func (c *Coordinator) awaitClients(ctx context.Context) (network.Client, network.Client, error) {
	p1Client, err := c.clientProvider.NextClient(ctx)
	if err != nil {
		return nil, nil, errors.Wrap(err, "Error when getting a Client")
	}
	p1Client.Start()

	for {
		p2Client, err := c.clientProvider.NextClient(ctx)
		if err != nil {
			p1Client.Disconnect()
			return nil, nil, errors.Wrap(err, "Error when getting a Client")
		}
		p2Client.Start()

		select {
		case <-p1Client.DisconnectCh():
			// p1 disconnected while waiting for p2. Set p1=p2
			// and find a new p2.
			p1Client = p2Client
		default:
			return p1Client, p2Client, nil
		}
	}
}

// run runs the main "loop" of the game coordinator. The loop waits for
// two players, creates a new games for these players, and wait for another
// pair of players. This method blocks until the context is cancelled or
// an error occurs. Always returns a non-nil error.
func (c *Coordinator) run(ctx context.Context) error {
	for {
		p1Client, p2Client, err := c.awaitClients(ctx)
		if err != nil {
			return errors.Wrap(err, "Error when awaiting clients")
		}

		c.logEntry.Debug("Starting a new game")
		err = c.startGame(ctx, p1Client, p2Client)
		if err != nil {
			// As we still own the clients here, make sure we stop them
			// before quiting ourselves
			p1Client.Disconnect()
			p2Client.Disconnect()
			return errors.Wrap(err, "Error starting game")
		}
	}
}

// Starts a new game for the two clients. The game is run on a new goroutine.
// This method blocks until the game has been created, but not until it has
// finished running.
func (c *Coordinator) startGame(ctx context.Context, p1Client, p2Client network.Client) error {
	game, err := model.CreateBasicGame()
	if err != nil {
		return errors.Wrap(err, "Error creating game")
	}
	ctrl, err := newController(c.logEntry.Logger, game, p1Client, p2Client)
	if err != nil {
		return errors.Wrap(err, "Error creating game controller")
	}

	// We create a new context that has a limited lifetime, so a game
	// cannot run forever.
	gameCtx, cancel := context.WithTimeout(ctx, GameMaxDuration)
	c.gamesWG.Add(1)
	go func() {
		defer c.gamesWG.Done()
		err := ctrl.Run(gameCtx)
		if err != nil && err != context.Canceled {
			c.logEntry.Errorf("Game stopped with an error: %+v", err)
		}
		cancel()
		// Disconnect the clients after the game is over
		p1Client.Disconnect()
		p2Client.Disconnect()
	}()
	return nil
}
<commit_msg>Fixes errors for game context cancellation<commit_after>package game

import (
	"context"
	"github.com/Sirupsen/logrus"
	"github.com/pkg/errors"
	"github.com/verath/archipelago/lib/common"
	"github.com/verath/archipelago/lib/game/model"
	"github.com/verath/archipelago/lib/network"
	"sync"
	"time"
)

const (
	// A timeout for the maximum length of a game before it is force-closed.
	// Used so that a client leaving a game open will not keep its resources
	// allocated forever.
	GameMaxDuration time.Duration = 45 * time.Minute
)

// The game coordinator is responsible for connecting players to
// a game. Once enough players has been found so that a game can
// be created, the game coordinator creates and starts a new game
// with the given players.
//
// The lifetime of the game coordinator is not tied to a single
// game but rather the entire lifetime of the game server.
type Coordinator struct {
	logEntry       *logrus.Entry
	clientProvider ClientProvider
	// WaitGroup for games created by the Coordinator
	gamesWG sync.WaitGroup
}

func NewCoordinator(log *logrus.Logger, clientProvider ClientProvider) (*Coordinator, error) {
	logEntry := common.ModuleLogEntry(log, "game/coordinator")
	return &Coordinator{
		logEntry:       logEntry,
		clientProvider: clientProvider,
	}, nil
}

// Starts and runs the Coordinator. This method blocks until the context is cancelled or
// an error occurs, and always returns a non-nil error.
func (c *Coordinator) Run(ctx context.Context) error {
	c.logEntry.Info("Starting")
	err := c.run(ctx)
	c.logEntry.Debug("Waiting for games to stop...")
	c.gamesWG.Wait()
	c.logEntry.Info("Stopped")
	return err
}

// Waits for two player connections to be made. If successful, the methods
// returns two started clients. These clients *must* be stopped. If the
// method returns an error the clients can be assumed to be stopped.
func (c *Coordinator) awaitClients(ctx context.Context) (network.Client, network.Client, error) {
	p1Client, err := c.clientProvider.NextClient(ctx)
	if err != nil {
		return nil, nil, errors.Wrap(err, "Error when getting a Client")
	}
	p1Client.Start()

	for {
		p2Client, err := c.clientProvider.NextClient(ctx)
		if err != nil {
			p1Client.Disconnect()
			return nil, nil, errors.Wrap(err, "Error when getting a Client")
		}
		p2Client.Start()

		select {
		case <-p1Client.DisconnectCh():
			// p1 disconnected while waiting for p2. Set p1=p2
			// and find a new p2.
			p1Client = p2Client
		default:
			return p1Client, p2Client, nil
		}
	}
}

// run runs the main "loop" of the game coordinator. The loop waits for
// two players, creates a new games for these players, and wait for another
// pair of players. This method blocks until the context is cancelled or
// an error occurs. Always returns a non-nil error.
func (c *Coordinator) run(ctx context.Context) error {
	for {
		p1Client, p2Client, err := c.awaitClients(ctx)
		if err != nil {
			return errors.Wrap(err, "Error when awaiting clients")
		}

		c.logEntry.Debug("Starting a new game")
		err = c.startGame(ctx, p1Client, p2Client)
		if err != nil {
			// As we still own the clients here, make sure we stop them
			// before quiting ourselves
			p1Client.Disconnect()
			p2Client.Disconnect()
			return errors.Wrap(err, "Error starting game")
		}
	}
}

// Starts a new game for the two clients. The game is run on a new goroutine.
// This method blocks until the game has been created, but not until it has
// finished running.
func (c *Coordinator) startGame(ctx context.Context, p1Client, p2Client network.Client) error {
	game, err := model.CreateBasicGame()
	if err != nil {
		return errors.Wrap(err, "Error creating game")
	}
	ctrl, err := newController(c.logEntry.Logger, game, p1Client, p2Client)
	if err != nil {
		return errors.Wrap(err, "Error creating game controller")
	}

	// We create a new context that has a limited lifetime, so a game
	// cannot run forever.
	gameCtx, cancel := context.WithTimeout(ctx, GameMaxDuration)
	c.gamesWG.Add(1)
	go func() {
		defer c.gamesWG.Done()
		err := ctrl.Run(gameCtx)
		if err != nil && errors.Cause(err) != context.Canceled {
			c.logEntry.Errorf("Game stopped with an error: %+v", err)
		}
		cancel()
		// Disconnect the clients after the game is over
		p1Client.Disconnect()
		p2Client.Disconnect()
	}()
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/dsvensson/go-xmmsclient/xmmsclient"
	"time"
)

func RepeatSomething(c *xmmsclient.Client) {
	for {
		fmt.Println(c.MainListPlugins(0))
		time.Sleep(time.Millisecond * 500)
	}
}

func main() {
	client := xmmsclient.NewClient("hello-from-go")

	client.Dial("localhost:xmms2")

	go RepeatSomething(client)

	fmt.Println("Plugins:")

	fmt.Println(client.MainListPlugins(0))

	select {}
}
<commit_msg>Update test client for shutdowns, and less verbosity.<commit_after>package main

import (
	"fmt"
	"github.com/dsvensson/go-xmmsclient/xmmsclient"
	"time"
)

func repeat(c *xmmsclient.Client) {
	for {
		value, err := c.PlaylistListEntries("_active")
		if err != nil {
			return
		}
		fmt.Println("repeat():", value)
		time.Sleep(time.Millisecond * 500)
	}
}

func main() {
	client := xmmsclient.NewClient("hello-from-go")

	client.Dial("localhost:xmms2")

	go repeat(client)

	value, err := client.PlaylistListEntries("_active")
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Println("  main():", value)

	time.Sleep(time.Second * 2)
	client.Close()
	time.Sleep(time.Second * 1)
}
<|endoftext|>


<commit_before>package main

import (
	"archive/zip"
	"bufio"
	"compress/flate"
	"encoding/binary"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
)

// from archive/zip struct.go

const (
	fileHeaderSignature      = 0x04034b50
	directoryHeaderSignature = 0x02014b50
	dataDescriptorSignature  = 0x08074b50 // de-facto standard; required by OS X Finder
	fileHeaderLen            = 30         // + filename + extra
	dataDescriptorLen        = 16         // four uint32: descriptor signature, crc32, compressed size, size
	dataDescriptor64Len      = 24         // descriptor with 8 byte sizes

	// version numbers
	zipVersion20 = 20 // 2.0
	zipVersion45 = 45 // 4.5 (reads and writes zip64 archives)
)

type readBuf []byte

func (b *readBuf) uint16() uint16 {
	v := binary.LittleEndian.Uint16(*b)
	*b = (*b)[2:]
	return v
}

func (b *readBuf) uint32() uint32 {
	v := binary.LittleEndian.Uint32(*b)
	*b = (*b)[4:]
	return v
}

func (b *readBuf) uint64() uint64 {
	v := binary.LittleEndian.Uint64(*b)
	*b = (*b)[8:]
	return v
}

func main() {
	debug := flag.Bool("debug", false, "print debug info")
	view := flag.Bool("v", false, "view list")

	flag.Parse()

	if flag.NArg() != 1 {
		log.Fatal("usage: ", path.Base(os.Args[0]), " {zip-file}")
	}

	zipfile := flag.Arg(0)

	f, err := os.Open(zipfile)
	if err != nil {
		log.Fatal("open ", err)
	}

	defer f.Close()

	r := bufio.NewReader(f)

	for {
		var fh [fileHeaderLen]byte

		if _, err := io.ReadFull(r, fh[:]); err != nil {
			log.Fatal("file header ", err)
		}

		if *debug {
			fmt.Println(hex.Dump(fh[:]))
		}

		b := readBuf(fh[:])
		magic := b.uint32()
		version := b.uint16()
		flags := b.uint16()
		comp := b.uint16()
		ctime := b.uint16()
		cdate := b.uint16()
		crc32 := b.uint32()
		clen := b.uint32()
		ulen := b.uint32()
		flen := b.uint16()
		elen := b.uint16()

		ctype := ""

		if magic == directoryHeaderSignature {
			// got central directory. Done
			log.Println("found central directory")
			break
		}

		if magic != fileHeaderSignature {
			log.Fatal("invalid file header signature ", fmt.Sprintf("%08x", magic))
		}

		if *debug {
			fmt.Println()
			fmt.Printf("magic   %08x\n", magic)
			fmt.Printf("version %04x\n", version)
			fmt.Printf("flags   %04x\n", flags)
			fmt.Printf("comp    %04x\n", comp)
			fmt.Printf("time    %04x\n", ctime)
			fmt.Printf("date    %04x\n", cdate)
			fmt.Printf("crc32   %08x\n", crc32)
			fmt.Printf("compressed size   %d\n", clen)
			fmt.Printf("uncompressed size %d\n", ulen)
			fmt.Printf("filename length   %d\n", flen)
			fmt.Printf("extra length      %d\n", elen)
		}

		fn := make([]byte, flen)
		if _, err := io.ReadFull(r, fn); err != nil {
			log.Fatal("read file name ", err)
		}

		if *debug {
			fmt.Println()
			fmt.Println("filename", string(fn))
		}

		if elen > 0 {
			if _, err := io.CopyN(ioutil.Discard, r, int64(elen)); err != nil {
				log.Fatal("read extra ", err)
			}
		}

		switch comp {
		case zip.Deflate:
			ctype = "Defl:N"

			w := ioutil.Discard
			if !*view {
				filename := string(fn)
				fmt.Println("inflating:", filename)

				dir := filepath.Dir(filename)
				if dir != "" {
					if err := os.MkdirAll(dir, 0755); err != nil {
						log.Println("mkdir", dir, err)
					}
				}

				if f, err := os.Create(filename); err != nil {
					log.Fatal("create ", fn, err)
				} else {
					w = f
				}
			}

			dec := flate.NewReader(r)
			n, err := io.Copy(w, dec)
			if *debug {
				fmt.Println("decoded", n, "bytes")
			}
			if err != nil {
				if wc, ok := w.(io.Closer); ok {
					wc.Close()
					os.Remove(string(fn))
				}

				log.Fatal("decode file ", err)
			} else {
				dec.Close()

				if wc, ok := w.(io.Closer); ok {
					wc.Close()
				}
			}

		case zip.Store:
			ctype = "Stored"

			if ulen > 0 {
				n, err := io.CopyN(ioutil.Discard, r, int64(ulen))
				if *debug {
					fmt.Println("read", n, "bytes")
				}
				if err != nil {
					log.Fatal("read file ", err)
				}
			} else {
				log.Fatal("missing lenght")
			}

		default:
			log.Fatal("unsupported compression mode ", comp)
		}

		if (flags & 0x08) != 0 {
			// data descriptor
			var dd [dataDescriptorLen]byte

			if _, err := io.ReadFull(r, dd[:]); err != nil {
				log.Fatal("data descriptor", err)
			}

			b := readBuf(dd[:])
			magic := b.uint32()
			crc32 = b.uint32()
			clen = b.uint32()
			ulen = b.uint32()

			if magic != dataDescriptorSignature {
				log.Fatal("invalid data descriptor signature ", magic)
			}

			if *debug {
				fmt.Println()
				fmt.Printf("magic   %08x\n", magic)
				fmt.Printf("crc32   %08x\n", crc32)
				fmt.Printf("compressed size   %d\n", clen)
				fmt.Printf("uncompressed size %d\n", ulen)
			}
		}

		if *view {
			pc := 100 - (clen * 100 / ulen)
			fmt.Printf("%8d  %6s  %8d  %2d%%  %08x  %s\n", ulen, ctype, clen, pc, crc32, fn)
		}
	}
}
<commit_msg>Added -out option to write restored files in output zip file<commit_after>package main

import (
	"archive/zip"
	"bufio"
	"compress/flate"
	"encoding/binary"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path"
	"path/filepath"
)

// from archive/zip struct.go

const (
	fileHeaderSignature      = 0x04034b50
	directoryHeaderSignature = 0x02014b50
	dataDescriptorSignature  = 0x08074b50 // de-facto standard; required by OS X Finder
	fileHeaderLen            = 30         // + filename + extra
	dataDescriptorLen        = 16         // four uint32: descriptor signature, crc32, compressed size, size
	dataDescriptor64Len      = 24         // descriptor with 8 byte sizes

	// version numbers
	zipVersion20 = 20 // 2.0
	zipVersion45 = 45 // 4.5 (reads and writes zip64 archives)
)

type readBuf []byte

func (b *readBuf) uint16() uint16 {
	v := binary.LittleEndian.Uint16(*b)
	*b = (*b)[2:]
	return v
}

func (b *readBuf) uint32() uint32 {
	v := binary.LittleEndian.Uint32(*b)
	*b = (*b)[4:]
	return v
}

func (b *readBuf) uint64() uint64 {
	v := binary.LittleEndian.Uint64(*b)
	*b = (*b)[8:]
	return v
}

func main() {
	debug := flag.Bool("debug", false, "print debug info")
	view := flag.Bool("v", false, "view list")
	out := flag.String("out", "", "write recovered files to output zip file")

	flag.Parse()

	if flag.NArg() != 1 {
		log.Fatal("usage: ", path.Base(os.Args[0]), " {zip-file}")
	}

	zipfile := flag.Arg(0)

	f, err := os.Open(zipfile)
	if err != nil {
		log.Fatal("open ", err)
	}

	defer f.Close()

	r := bufio.NewReader(f)

	var outz *zip.Writer

	if len(*out) > 0 {
		outf, err := os.Create(*out)
		if err != nil {
			log.Fatal("create output", err)
		}

		outz = zip.NewWriter(outf)

		defer func() {
			outz.Close()
			outf.Close()
		}()
	}

Loop:
	for {
		var fh [fileHeaderLen]byte

		if _, err := io.ReadFull(r, fh[:]); err != nil {
			log.Println("file header", err)
			break Loop
		}

		if *debug {
			fmt.Println(hex.Dump(fh[:]))
		}

		b := readBuf(fh[:])
		magic := b.uint32()
		version := b.uint16()
		flags := b.uint16()
		comp := b.uint16()
		ctime := b.uint16()
		cdate := b.uint16()
		crc32 := b.uint32()
		clen := b.uint32()
		ulen := b.uint32()
		flen := b.uint16()
		elen := b.uint16()

		ctype := ""

		if magic == directoryHeaderSignature {
			// got central directory. Done
			log.Println("found central directory")
			break Loop
		}

		if magic != fileHeaderSignature {
			log.Fatal("invalid file header signature ", fmt.Sprintf("%08x", magic))
		}

		if *debug {
			fmt.Println()
			fmt.Printf("magic   %08x\n", magic)
			fmt.Printf("version %04x\n", version)
			fmt.Printf("flags   %04x\n", flags)
			fmt.Printf("comp    %04x\n", comp)
			fmt.Printf("time    %04x\n", ctime)
			fmt.Printf("date    %04x\n", cdate)
			fmt.Printf("crc32   %08x\n", crc32)
			fmt.Printf("compressed size   %d\n", clen)
			fmt.Printf("uncompressed size %d\n", ulen)
			fmt.Printf("filename length   %d\n", flen)
			fmt.Printf("extra length      %d\n", elen)
		}

		fn := make([]byte, flen)
		if _, err := io.ReadFull(r, fn); err != nil {
			log.Println("read file name", err)
			break Loop
		}

		if *debug {
			fmt.Println()
			fmt.Println("filename", string(fn))
		}

		if elen > 0 {
			if _, err := io.CopyN(ioutil.Discard, r, int64(elen)); err != nil {
				log.Println("read extra", err)
				break Loop
			}
		}

		filename := string(fn)

		switch comp {
		case zip.Deflate:
			ctype = "Defl:N"

			var w io.Writer

			if *view {
				w = ioutil.Discard
			} else if outz != nil {
				fmt.Println("adding:", filename)
				if f, err := outz.Create(filename); err != nil {
					log.Fatal("create zip entry ", filename, err)
				} else {
					w = f
				}
			} else {
				fmt.Println("inflating:", filename)

				dir := filepath.Dir(filename)
				if dir != "" {
					if err := os.MkdirAll(dir, 0755); err != nil {
						log.Println("mkdir", dir, err)
					}
				}

				if f, err := os.Create(filename); err != nil {
					log.Fatal("create ", filename, err)
				} else {
					w = f
				}
			}

			dec := flate.NewReader(r)
			n, err := io.Copy(w, dec)
			if *debug {
				fmt.Println("decoded", n, "bytes")
			}
			if err != nil {
				if wc, ok := w.(io.Closer); ok {
					wc.Close()
					os.Remove(filename)
				}

				log.Println("decode file", err)
				break Loop
			} else {
				dec.Close()

				if wc, ok := w.(io.Closer); ok {
					wc.Close()
				}
			}

		case zip.Store:
			ctype = "Stored"

			if ulen > 0 {
				n, err := io.CopyN(ioutil.Discard, r, int64(ulen))
				if *debug {
					fmt.Println("read", n, "bytes")
				}
				if err != nil {
					log.Fatal("read file ", err)
				}
			} else {
				log.Fatal("missing lenght")
			}

		default:
			log.Fatal("unsupported compression mode ", comp)
		}

		if (flags & 0x08) != 0 {
			// data descriptor
			var dd [dataDescriptorLen]byte

			if _, err := io.ReadFull(r, dd[:]); err != nil {
				log.Fatal("data descriptor", err)
			}

			b := readBuf(dd[:])
			magic := b.uint32()
			crc32 = b.uint32()
			clen = b.uint32()
			ulen = b.uint32()

			if magic != dataDescriptorSignature {
				log.Fatal("invalid data descriptor signature ", magic)
			}

			if *debug {
				fmt.Println()
				fmt.Printf("magic   %08x\n", magic)
				fmt.Printf("crc32   %08x\n", crc32)
				fmt.Printf("compressed size   %d\n", clen)
				fmt.Printf("uncompressed size %d\n", ulen)
			}
		}

		if *view {
			pc := 100 - (clen * 100 / ulen)
			fmt.Printf("%8d  %6s  %8d  %2d%%  %08x  %s\n", ulen, ctype, clen, pc, crc32, filename)
		}
	}
}
<|endoftext|>


<commit_before>// Generate Plugins is a small program that updates the lists of plugins in
// command/internal_plugin_list.go so they will be compiled into the main
// terraform binary.
package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

const target = "command/internal_plugin_list.go"

func main() {
	wd, _ := os.Getwd()
	if filepath.Base(wd) != "terraform" {
		log.Fatalf("This program must be invoked in the terraform project root; in %s", wd)
	}

	// Collect all of the data we need about plugins we have in the project
	providers, err := discoverProviders()
	if err != nil {
		log.Fatalf("Failed to discover providers: %s", err)
	}

	provisioners, err := discoverProvisioners()
	if err != nil {
		log.Fatalf("Failed to discover provisioners: %s", err)
	}

	// Do some simple code generation and templating
	output := source
	output = strings.Replace(output, "IMPORTS", makeImports(providers, provisioners), 1)
	output = strings.Replace(output, "PROVIDERS", makeProviderMap(providers), 1)
	output = strings.Replace(output, "PROVISIONERS", makeProvisionerMap(provisioners), 1)

	// TODO sort the lists of plugins so we are not subjected to random OS ordering of the plugin lists

	// Write our generated code to the command/plugin.go file
	file, err := os.Create(target)
	defer file.Close()
	if err != nil {
		log.Fatalf("Failed to open %s for writing: %s", target, err)
	}

	_, err = file.WriteString(output)
	if err != nil {
		log.Fatalf("Failed writing to %s: %s", target, err)
	}

	log.Printf("Generated %s", target)
}

type plugin struct {
	Package    string // Package name from ast  remoteexec
	PluginName string // Path via deriveName()  remote-exec
	TypeName   string // Type of plugin         provisioner
	Path       string // Relative import path   builtin/provisioners/remote-exec
	ImportName string // See deriveImport()     remoteexecprovisioner
}

// makeProviderMap creates a map of providers like this:
//
// var InternalProviders = map[string]plugin.ProviderFunc{
// 	"aws":        aws.Provider,
// 	"azurerm":    azurerm.Provider,
// 	"cloudflare": cloudflare.Provider,
func makeProviderMap(items []plugin) string {
	output := ""
	for _, item := range items {
		output += fmt.Sprintf("\t\"%s\":   %s.%s,\n", item.PluginName, item.ImportName, item.TypeName)
	}
	return output
}

// makeProvisionerMap creates a map of provisioners like this:
//
//	"file":        func() terraform.ResourceProvisioner { return new(file.ResourceProvisioner) },
//	"local-exec":  func() terraform.ResourceProvisioner { return new(localexec.ResourceProvisioner) },
//	"remote-exec": func() terraform.ResourceProvisioner { return new(remoteexec.ResourceProvisioner) },
//
// This is more verbose than the Provider case because there is no corresponding
// Provisioner function.
func makeProvisionerMap(items []plugin) string {
	output := ""
	for _, item := range items {
		output += fmt.Sprintf("\t\"%s\": func() terraform.ResourceProvisioner { return new(%s.%s) },\n", item.PluginName, item.ImportName, item.TypeName)
	}
	return output
}

func makeImports(providers, provisioners []plugin) string {
	plugins := []string{}

	for _, provider := range providers {
		plugins = append(plugins, fmt.Sprintf("\t%s \"github.com/hashicorp/terraform/%s\"\n", provider.ImportName, filepath.ToSlash(provider.Path)))
	}

	for _, provisioner := range provisioners {
		plugins = append(plugins, fmt.Sprintf("\t%s \"github.com/hashicorp/terraform/%s\"\n", provisioner.ImportName, filepath.ToSlash(provisioner.Path)))
	}

	// Make things pretty
	sort.Strings(plugins)

	return strings.Join(plugins, "")
}

// listDirectories recursively lists directories under the specified path
func listDirectories(path string) ([]string, error) {
	names := []string{}
	items, err := ioutil.ReadDir(path)
	if err != nil {
		return names, err
	}

	for _, item := range items {
		// We only want directories
		if item.IsDir() {
			if item.Name() == "test-fixtures" {
				continue
			}
			currentDir := filepath.Join(path, item.Name())
			names = append(names, currentDir)

			// Do some recursion
			subNames, err := listDirectories(currentDir)
			if err == nil {
				names = append(names, subNames...)
			}
		}
	}

	return names, nil
}

// deriveName determines the name of the plugin relative to the specified root
// path.
func deriveName(root, full string) string {
	short, _ := filepath.Rel(root, full)
	bits := strings.Split(short, string(os.PathSeparator))
	return strings.Join(bits, "-")
}

// deriveImport will build a unique import identifier based on packageName and
// the result of deriveName(). This is important for disambigutating between
// providers and provisioners that have the same name. This will be something
// like:
//
//	remote-exec -> remoteexecprovisioner
//
// which is long, but is deterministic and unique.
func deriveImport(typeName, derivedName string) string {
	return strings.Replace(derivedName, "-", "", -1) + strings.ToLower(typeName)
}

// discoverTypesInPath searches for types of typeID in path using go's ast and
// returns a list of plugins it finds.
func discoverTypesInPath(path, typeID, typeName string) ([]plugin, error) {
	pluginTypes := []plugin{}

	dirs, err := listDirectories(path)
	if err != nil {
		return pluginTypes, err
	}

	for _, dir := range dirs {
		fset := token.NewFileSet()
		goPackages, err := parser.ParseDir(fset, dir, nil, parser.AllErrors)
		if err != nil {
			return pluginTypes, fmt.Errorf("Failed parsing directory %s: %s", dir, err)
		}

		for _, goPackage := range goPackages {
			ast.PackageExports(goPackage)
			ast.Inspect(goPackage, func(n ast.Node) bool {
				switch x := n.(type) {
				case *ast.FuncDecl:
					// If we get a function then we will check the function name
					// against typeName and the function return type (Results)
					// against typeID.
					//
					// There may be more than one return type but in the target
					// case there should only be one. Also the return type is a
					// ast.SelectorExpr which means we have multiple nodes.
					// We'll read all of them as ast.Ident (identifier), join
					// them via . to get a string like terraform.ResourceProvider
					// and see if it matches our expected typeID
					//
					// This is somewhat verbose but prevents us from identifying
					// the wrong types if the function name is amiguous or if
					// there are other subfolders added later.
					if x.Name.Name == typeName && len(x.Type.Results.List) == 1 {
						node := x.Type.Results.List[0].Type
						typeIdentifiers := []string{}
						ast.Inspect(node, func(m ast.Node) bool {
							switch y := m.(type) {
							case *ast.Ident:
								typeIdentifiers = append(typeIdentifiers, y.Name)
							}
							// We need all of the identifiers to join so we
							// can't break early here.
							return true
						})
						if strings.Join(typeIdentifiers, ".") == typeID {
							derivedName := deriveName(path, dir)
							pluginTypes = append(pluginTypes, plugin{
								Package:    goPackage.Name,
								PluginName: derivedName,
								ImportName: deriveImport(x.Name.Name, derivedName),
								TypeName:   x.Name.Name,
								Path:       dir,
							})
						}
					}
				case *ast.TypeSpec:
					// In the simpler case we will simply check whether the type
					// declaration has the name we were looking for.
					if x.Name.Name == typeID {
						derivedName := deriveName(path, dir)
						pluginTypes = append(pluginTypes, plugin{
							Package:    goPackage.Name,
							PluginName: derivedName,
							ImportName: deriveImport(x.Name.Name, derivedName),
							TypeName:   x.Name.Name,
							Path:       dir,
						})
						// The AST stops parsing when we return false. Once we
						// find the symbol we want we can stop parsing.
						return false
					}
				}
				return true
			})
		}
	}

	return pluginTypes, nil
}

func discoverProviders() ([]plugin, error) {
	path := "./builtin/providers"
	typeID := "terraform.ResourceProvider"
	typeName := "Provider"
	return discoverTypesInPath(path, typeID, typeName)
}

func discoverProvisioners() ([]plugin, error) {
	path := "./builtin/provisioners"
	typeID := "ResourceProvisioner"
	typeName := ""
	return discoverTypesInPath(path, typeID, typeName)
}

const source = `// +build !core

//
// This file is automatically generated by scripts/generate-plugins.go -- Do not edit!
//
package command

import (
IMPORTS
	"github.com/hashicorp/terraform/plugin"
	"github.com/hashicorp/terraform/terraform"
)

var InternalProviders = map[string]plugin.ProviderFunc{
PROVIDERS
}

var InternalProvisioners = map[string]plugin.ProvisionerFunc{
PROVISIONERS
}

`
<commit_msg>scripts: update internal plugin gen to support new provisioner<commit_after>// Generate Plugins is a small program that updates the lists of plugins in
// command/internal_plugin_list.go so they will be compiled into the main
// terraform binary.
package main

import (
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

const target = "command/internal_plugin_list.go"

func main() {
	wd, _ := os.Getwd()
	if filepath.Base(wd) != "terraform" {
		log.Fatalf("This program must be invoked in the terraform project root; in %s", wd)
	}

	// Collect all of the data we need about plugins we have in the project
	providers, err := discoverProviders()
	if err != nil {
		log.Fatalf("Failed to discover providers: %s", err)
	}

	provisioners, err := discoverProvisioners()
	if err != nil {
		log.Fatalf("Failed to discover provisioners: %s", err)
	}

	// Do some simple code generation and templating
	output := source
	output = strings.Replace(output, "IMPORTS", makeImports(providers, provisioners), 1)
	output = strings.Replace(output, "PROVIDERS", makeProviderMap(providers), 1)
	output = strings.Replace(output, "PROVISIONERS", makeProvisionerMap(provisioners), 1)

	// TODO sort the lists of plugins so we are not subjected to random OS ordering of the plugin lists

	// Write our generated code to the command/plugin.go file
	file, err := os.Create(target)
	defer file.Close()
	if err != nil {
		log.Fatalf("Failed to open %s for writing: %s", target, err)
	}

	_, err = file.WriteString(output)
	if err != nil {
		log.Fatalf("Failed writing to %s: %s", target, err)
	}

	log.Printf("Generated %s", target)
}

type plugin struct {
	Package    string // Package name from ast  remoteexec
	PluginName string // Path via deriveName()  remote-exec
	TypeName   string // Type of plugin         provisioner
	Path       string // Relative import path   builtin/provisioners/remote-exec
	ImportName string // See deriveImport()     remoteexecprovisioner
}

// makeProviderMap creates a map of providers like this:
//
// var InternalProviders = map[string]plugin.ProviderFunc{
// 	"aws":        aws.Provider,
// 	"azurerm":    azurerm.Provider,
// 	"cloudflare": cloudflare.Provider,
func makeProviderMap(items []plugin) string {
	output := ""
	for _, item := range items {
		output += fmt.Sprintf("\t\"%s\":   %s.%s,\n", item.PluginName, item.ImportName, item.TypeName)
	}
	return output
}

// makeProvisionerMap creates a map of provisioners like this:
//
//	"file":        func() terraform.ResourceProvisioner { return new(file.ResourceProvisioner) },
//	"local-exec":  func() terraform.ResourceProvisioner { return new(localexec.ResourceProvisioner) },
//	"remote-exec": func() terraform.ResourceProvisioner { return new(remoteexec.ResourceProvisioner) },
//
// This is more verbose than the Provider case because there is no corresponding
// Provisioner function.
func makeProvisionerMap(items []plugin) string {
	output := ""
	for _, item := range items {
		output += fmt.Sprintf("\t\"%s\":   %s.%s,\n", item.PluginName, item.ImportName, item.TypeName)
	}
	return output
}

func makeImports(providers, provisioners []plugin) string {
	plugins := []string{}

	for _, provider := range providers {
		plugins = append(plugins, fmt.Sprintf("\t%s \"github.com/hashicorp/terraform/%s\"\n", provider.ImportName, filepath.ToSlash(provider.Path)))
	}

	for _, provisioner := range provisioners {
		plugins = append(plugins, fmt.Sprintf("\t%s \"github.com/hashicorp/terraform/%s\"\n", provisioner.ImportName, filepath.ToSlash(provisioner.Path)))
	}

	// Make things pretty
	sort.Strings(plugins)

	return strings.Join(plugins, "")
}

// listDirectories recursively lists directories under the specified path
func listDirectories(path string) ([]string, error) {
	names := []string{}
	items, err := ioutil.ReadDir(path)
	if err != nil {
		return names, err
	}

	for _, item := range items {
		// We only want directories
		if item.IsDir() {
			if item.Name() == "test-fixtures" {
				continue
			}
			currentDir := filepath.Join(path, item.Name())
			names = append(names, currentDir)

			// Do some recursion
			subNames, err := listDirectories(currentDir)
			if err == nil {
				names = append(names, subNames...)
			}
		}
	}

	return names, nil
}

// deriveName determines the name of the plugin relative to the specified root
// path.
func deriveName(root, full string) string {
	short, _ := filepath.Rel(root, full)
	bits := strings.Split(short, string(os.PathSeparator))
	return strings.Join(bits, "-")
}

// deriveImport will build a unique import identifier based on packageName and
// the result of deriveName(). This is important for disambigutating between
// providers and provisioners that have the same name. This will be something
// like:
//
//	remote-exec -> remoteexecprovisioner
//
// which is long, but is deterministic and unique.
func deriveImport(typeName, derivedName string) string {
	return strings.Replace(derivedName, "-", "", -1) + strings.ToLower(typeName)
}

// discoverTypesInPath searches for types of typeID in path using go's ast and
// returns a list of plugins it finds.
func discoverTypesInPath(path, typeID, typeName string) ([]plugin, error) {
	pluginTypes := []plugin{}

	dirs, err := listDirectories(path)
	if err != nil {
		return pluginTypes, err
	}

	for _, dir := range dirs {
		fset := token.NewFileSet()
		goPackages, err := parser.ParseDir(fset, dir, nil, parser.AllErrors)
		if err != nil {
			return pluginTypes, fmt.Errorf("Failed parsing directory %s: %s", dir, err)
		}

		for _, goPackage := range goPackages {
			ast.PackageExports(goPackage)
			ast.Inspect(goPackage, func(n ast.Node) bool {
				switch x := n.(type) {
				case *ast.FuncDecl:
					// If we get a function then we will check the function name
					// against typeName and the function return type (Results)
					// against typeID.
					//
					// There may be more than one return type but in the target
					// case there should only be one. Also the return type is a
					// ast.SelectorExpr which means we have multiple nodes.
					// We'll read all of them as ast.Ident (identifier), join
					// them via . to get a string like terraform.ResourceProvider
					// and see if it matches our expected typeID
					//
					// This is somewhat verbose but prevents us from identifying
					// the wrong types if the function name is amiguous or if
					// there are other subfolders added later.
					if x.Name.Name == typeName && len(x.Type.Results.List) == 1 {
						node := x.Type.Results.List[0].Type
						typeIdentifiers := []string{}
						ast.Inspect(node, func(m ast.Node) bool {
							switch y := m.(type) {
							case *ast.Ident:
								typeIdentifiers = append(typeIdentifiers, y.Name)
							}
							// We need all of the identifiers to join so we
							// can't break early here.
							return true
						})
						if strings.Join(typeIdentifiers, ".") == typeID {
							derivedName := deriveName(path, dir)
							pluginTypes = append(pluginTypes, plugin{
								Package:    goPackage.Name,
								PluginName: derivedName,
								ImportName: deriveImport(x.Name.Name, derivedName),
								TypeName:   x.Name.Name,
								Path:       dir,
							})
						}
					}
				case *ast.TypeSpec:
					// In the simpler case we will simply check whether the type
					// declaration has the name we were looking for.
					if x.Name.Name == typeID {
						derivedName := deriveName(path, dir)
						pluginTypes = append(pluginTypes, plugin{
							Package:    goPackage.Name,
							PluginName: derivedName,
							ImportName: deriveImport(x.Name.Name, derivedName),
							TypeName:   x.Name.Name,
							Path:       dir,
						})
						// The AST stops parsing when we return false. Once we
						// find the symbol we want we can stop parsing.
						return false
					}
				}
				return true
			})
		}
	}

	return pluginTypes, nil
}

func discoverProviders() ([]plugin, error) {
	path := "./builtin/providers"
	typeID := "terraform.ResourceProvider"
	typeName := "Provider"
	return discoverTypesInPath(path, typeID, typeName)
}

func discoverProvisioners() ([]plugin, error) {
	path := "./builtin/provisioners"
	typeID := "terraform.ResourceProvisioner"
	typeName := "Provisioner"
	return discoverTypesInPath(path, typeID, typeName)
}

const source = `// +build !core

//
// This file is automatically generated by scripts/generate-plugins.go -- Do not edit!
//
package command

import (
IMPORTS
	"github.com/hashicorp/terraform/plugin"
	"github.com/hashicorp/terraform/terraform"
)

var InternalProviders = map[string]plugin.ProviderFunc{
PROVIDERS
}

var InternalProvisioners = map[string]plugin.ProvisionerFunc{
PROVISIONERS
}

`
<|endoftext|>