<commit_before>package cli

import (
	"context"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/OJ/gobuster/v3/libgobuster"
)

const ruler = "==============================================================="
const cliProgressUpdate = 500 * time.Millisecond

func banner() {
	fmt.Printf("Gobuster v%s\n", libgobuster.VERSION)
	fmt.Println("by OJ Reeves (@TheColonial) & Christian Mehlmauer (@firefart)")
}

type outputType struct {
	Mu              *sync.RWMutex
	MaxCharsWritten int
}

// right pad a string
// nolint:unparam
func rightPad(s string, padStr string, overallLen int) string {
	strLen := len(s)
	if overallLen <= strLen {
		return s
	}

	toPad := overallLen - strLen - 1
	pad := strings.Repeat(padStr, toPad)
	return fmt.Sprintf("%s%s", s, pad)
}

// resultWorker outputs the results as they come in. This needs to be a range and should not handle
// the context so the channel always has a receiver and libgobuster will not block.
func resultWorker(g *libgobuster.Gobuster, filename string, wg *sync.WaitGroup, output *outputType) {
	defer wg.Done()

	var f *os.File
	var err error
	if filename != "" {
		f, err = os.Create(filename)
		if err != nil {
			g.LogError.Fatalf("error on creating output file: %v", err)
		}
		defer f.Close()
	}

	for r := range g.Results() {
		s, err := r.ToString(g)
		if err != nil {
			g.LogError.Fatal(err)
		}
		if s != "" {
			s = strings.TrimSpace(s)
			output.Mu.Lock()
			w, _ := fmt.Printf("\r%s\n", rightPad(s, " ", output.MaxCharsWritten))
			// -1 to remove the newline, otherwise it's always bigger
			if (w - 1) > output.MaxCharsWritten {
				output.MaxCharsWritten = w - 1
			}
			output.Mu.Unlock()
			if f != nil {
				err = writeToFile(f, s)
				if err != nil {
					g.LogError.Fatalf("error on writing output file: %v", err)
				}
			}
		}
	}
}

// errorWorker outputs the errors as they come in. This needs to be a range and should not handle
// the context so the channel always has a receiver and libgobuster will not block.
func errorWorker(g *libgobuster.Gobuster, wg *sync.WaitGroup, output *outputType) {
	defer wg.Done()

	for e := range g.Errors() {
		if !g.Opts.Quiet {
			output.Mu.Lock()
			g.LogError.Printf("[!] %v", e)
			output.Mu.Unlock()
		}
	}
}

// progressWorker outputs the progress every tick. It will stop once cancel() is called
// on the context
func progressWorker(c context.Context, g *libgobuster.Gobuster, wg *sync.WaitGroup, output *outputType) {
	defer wg.Done()

	tick := time.NewTicker(cliProgressUpdate)

	for {
		select {
		case <-tick.C:
			if !g.Opts.Quiet && !g.Opts.NoProgress {
				g.RequestsCountMutex.RLock()
				output.Mu.Lock()
				var charsWritten int
				if g.Opts.Wordlist == "-" {
					s := fmt.Sprintf("\rProgress: %d", g.RequestsIssued)
					s = rightPad(s, " ", output.MaxCharsWritten)
					charsWritten, _ = fmt.Fprint(os.Stderr, s)
					// only print status if we already read in the wordlist
				} else if g.RequestsExpected > 0 {
					s := fmt.Sprintf("\rProgress: %d / %d (%3.2f%%)", g.RequestsIssued, g.RequestsExpected, (float32(g.RequestsIssued) * 100.0 / float32(g.RequestsExpected)))
					s = rightPad(s, " ", output.MaxCharsWritten)
					charsWritten, _ = fmt.Fprint(os.Stderr, s)
				}
				if charsWritten > output.MaxCharsWritten {
					output.MaxCharsWritten = charsWritten
				}

				output.Mu.Unlock()
				g.RequestsCountMutex.RUnlock()
			}
		case <-c.Done():
			return
		}
	}
}

func writeToFile(f *os.File, output string) error {
	_, err := f.WriteString(fmt.Sprintf("%s\n", output))
	if err != nil {
		return fmt.Errorf("[!] Unable to write to file %w", err)
	}
	return nil
}

// Gobuster is the main entry point for the CLI
func Gobuster(prevCtx context.Context, opts *libgobuster.Options, plugin libgobuster.GobusterPlugin) error {
	// Sanity checks
	if opts == nil {
		return fmt.Errorf("please provide valid options")
	}

	if plugin == nil {
		return fmt.Errorf("please provide a valid plugin")
	}

	ctx, cancel := context.WithCancel(prevCtx)
	defer cancel()

	gobuster, err := libgobuster.NewGobuster(ctx, opts, plugin)
	if err != nil {
		return err
	}

	if !opts.Quiet {
		fmt.Println(ruler)
		banner()
		fmt.Println(ruler)
		c, err := gobuster.GetConfigString()
		if err != nil {
			return fmt.Errorf("error on creating config string: %w", err)
		}
		fmt.Println(c)
		fmt.Println(ruler)
		gobuster.LogInfo.Printf("Starting gobuster in %s mode", plugin.Name())
		fmt.Println(ruler)
	}

	// our waitgroup for all goroutines
	// this ensures all goroutines are finished
	// when we call wg.Wait()
	var wg sync.WaitGroup

	outputMutex := new(sync.RWMutex)
	o := &outputType{
		Mu:              outputMutex,
		MaxCharsWritten: 0,
	}

	wg.Add(1)
	go resultWorker(gobuster, opts.OutputFilename, &wg, o)

	if !opts.Quiet && !opts.NoError {
		wg.Add(1)
		go errorWorker(gobuster, &wg, o)
	}

	if !opts.Quiet && !opts.NoProgress {
		// if not quiet add a new workgroup entry and start the goroutine
		wg.Add(1)
		go progressWorker(ctx, gobuster, &wg, o)
	}

	err = gobuster.Start()

	// call cancel func so progressWorker will exit (the only goroutine in this
	// file using the context) and to free resources
	cancel()
	// wait for all spun up goroutines to finish (all have to call wg.Done())
	wg.Wait()

	// Late error checking to finish all threads
	if err != nil {
		return err
	}

	if !opts.Quiet {
		// clear stderr progress
		fmt.Fprintf(os.Stderr, "\r%s\n", rightPad("", " ", o.MaxCharsWritten))
		fmt.Println(ruler)
		gobuster.LogInfo.Println("Finished")
		fmt.Println(ruler)
	}
	return nil
}
<commit_msg>actually consume errors<commit_after>package cli

import (
	"context"
	"fmt"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/OJ/gobuster/v3/libgobuster"
)

const ruler = "==============================================================="
const cliProgressUpdate = 500 * time.Millisecond

func banner() {
	fmt.Printf("Gobuster v%s\n", libgobuster.VERSION)
	fmt.Println("by OJ Reeves (@TheColonial) & Christian Mehlmauer (@firefart)")
}

type outputType struct {
	Mu              *sync.RWMutex
	MaxCharsWritten int
}

// right pad a string
// nolint:unparam
func rightPad(s string, padStr string, overallLen int) string {
	strLen := len(s)
	if overallLen <= strLen {
		return s
	}

	toPad := overallLen - strLen - 1
	pad := strings.Repeat(padStr, toPad)
	return fmt.Sprintf("%s%s", s, pad)
}

// resultWorker outputs the results as they come in. This needs to be a range and should not handle
// the context so the channel always has a receiver and libgobuster will not block.
func resultWorker(g *libgobuster.Gobuster, filename string, wg *sync.WaitGroup, output *outputType) {
	defer wg.Done()

	var f *os.File
	var err error
	if filename != "" {
		f, err = os.Create(filename)
		if err != nil {
			g.LogError.Fatalf("error on creating output file: %v", err)
		}
		defer f.Close()
	}

	for r := range g.Results() {
		s, err := r.ToString(g)
		if err != nil {
			g.LogError.Fatal(err)
		}
		if s != "" {
			s = strings.TrimSpace(s)
			output.Mu.Lock()
			w, _ := fmt.Printf("\r%s\n", rightPad(s, " ", output.MaxCharsWritten))
			// -1 to remove the newline, otherwise it's always bigger
			if (w - 1) > output.MaxCharsWritten {
				output.MaxCharsWritten = w - 1
			}
			output.Mu.Unlock()
			if f != nil {
				err = writeToFile(f, s)
				if err != nil {
					g.LogError.Fatalf("error on writing output file: %v", err)
				}
			}
		}
	}
}

// errorWorker outputs the errors as they come in. This needs to be a range and should not handle
// the context so the channel always has a receiver and libgobuster will not block.
func errorWorker(g *libgobuster.Gobuster, wg *sync.WaitGroup, output *outputType) {
	defer wg.Done()

	for e := range g.Errors() {
		if !g.Opts.Quiet && !g.Opts.NoError {
			output.Mu.Lock()
			g.LogError.Printf("[!] %v", e)
			output.Mu.Unlock()
		}
	}
}

// progressWorker outputs the progress every tick. It will stop once cancel() is called
// on the context
func progressWorker(c context.Context, g *libgobuster.Gobuster, wg *sync.WaitGroup, output *outputType) {
	defer wg.Done()

	tick := time.NewTicker(cliProgressUpdate)

	for {
		select {
		case <-tick.C:
			if !g.Opts.Quiet && !g.Opts.NoProgress {
				g.RequestsCountMutex.RLock()
				output.Mu.Lock()
				var charsWritten int
				if g.Opts.Wordlist == "-" {
					s := fmt.Sprintf("\rProgress: %d", g.RequestsIssued)
					s = rightPad(s, " ", output.MaxCharsWritten)
					charsWritten, _ = fmt.Fprint(os.Stderr, s)
					// only print status if we already read in the wordlist
				} else if g.RequestsExpected > 0 {
					s := fmt.Sprintf("\rProgress: %d / %d (%3.2f%%)", g.RequestsIssued, g.RequestsExpected, (float32(g.RequestsIssued) * 100.0 / float32(g.RequestsExpected)))
					s = rightPad(s, " ", output.MaxCharsWritten)
					charsWritten, _ = fmt.Fprint(os.Stderr, s)
				}
				if charsWritten > output.MaxCharsWritten {
					output.MaxCharsWritten = charsWritten
				}

				output.Mu.Unlock()
				g.RequestsCountMutex.RUnlock()
			}
		case <-c.Done():
			return
		}
	}
}

func writeToFile(f *os.File, output string) error {
	_, err := f.WriteString(fmt.Sprintf("%s\n", output))
	if err != nil {
		return fmt.Errorf("[!] Unable to write to file %w", err)
	}
	return nil
}

// Gobuster is the main entry point for the CLI
func Gobuster(prevCtx context.Context, opts *libgobuster.Options, plugin libgobuster.GobusterPlugin) error {
	// Sanity checks
	if opts == nil {
		return fmt.Errorf("please provide valid options")
	}

	if plugin == nil {
		return fmt.Errorf("please provide a valid plugin")
	}

	ctx, cancel := context.WithCancel(prevCtx)
	defer cancel()

	gobuster, err := libgobuster.NewGobuster(ctx, opts, plugin)
	if err != nil {
		return err
	}

	if !opts.Quiet {
		fmt.Println(ruler)
		banner()
		fmt.Println(ruler)
		c, err := gobuster.GetConfigString()
		if err != nil {
			return fmt.Errorf("error on creating config string: %w", err)
		}
		fmt.Println(c)
		fmt.Println(ruler)
		gobuster.LogInfo.Printf("Starting gobuster in %s mode", plugin.Name())
		fmt.Println(ruler)
	}

	// our waitgroup for all goroutines
	// this ensures all goroutines are finished
	// when we call wg.Wait()
	var wg sync.WaitGroup

	outputMutex := new(sync.RWMutex)
	o := &outputType{
		Mu:              outputMutex,
		MaxCharsWritten: 0,
	}

	wg.Add(1)
	go resultWorker(gobuster, opts.OutputFilename, &wg, o)

	wg.Add(1)
	go errorWorker(gobuster, &wg, o)

	if !opts.Quiet && !opts.NoProgress {
		// if not quiet add a new workgroup entry and start the goroutine
		wg.Add(1)
		go progressWorker(ctx, gobuster, &wg, o)
	}

	err = gobuster.Start()

	// call cancel func so progressWorker will exit (the only goroutine in this
	// file using the context) and to free resources
	cancel()
	// wait for all spun up goroutines to finish (all have to call wg.Done())
	wg.Wait()

	// Late error checking to finish all threads
	if err != nil {
		return err
	}

	if !opts.Quiet {
		// clear stderr progress
		fmt.Fprintf(os.Stderr, "\r%s\n", rightPad("", " ", o.MaxCharsWritten))
		fmt.Println(ruler)
		gobuster.LogInfo.Println("Finished")
		fmt.Println(ruler)
	}
	return nil
}
<|endoftext|>


<commit_before>package xcodeutil

import (
	"bufio"
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	cmd "github.com/bitrise-io/xcode-test/command"
	log "github.com/bitrise-io/xcode-test/logutil"
	"github.com/bitrise-io/xcode-test/models"
	ver "github.com/hashicorp/go-version"
)

var (
	osRegexp          = regexp.MustCompile(`-- (.+) (\d\.\d) --`)
	deviceStateRegexp = regexp.MustCompile(` *(.+) \(([A-Z0-9-]+)\) \((.+)\)`)
)

//=======================================
// Utility
//=======================================

// a simulator info line should look like this:
//  iPhone 5s (EA1C7E48-8137-428C-A0A5-B2C63FF276EB) (Shutdown)
// or
//  iPhone 4s (51B10EBD-C949-49F5-A38B-E658F41640FF) (Shutdown) (unavailable, runtime profile not found)
func getSimInfoFromLine(lineStr string) (models.SimInfoModel, error) {
	baseInfosExp := regexp.MustCompile(`(?P<deviceName>[a-zA-Z].*[a-zA-Z0-9 -]*) \((?P<simulatorID>[a-zA-Z0-9-]{36})\) \((?P<status>[a-zA-Z]*)\)`)
	baseInfosRes := baseInfosExp.FindStringSubmatch(lineStr)
	if baseInfosRes == nil {
		return models.SimInfoModel{}, fmt.Errorf("No match found")
	}

	simInfo := models.SimInfoModel{
		Name:   baseInfosRes[1],
		SimID:  baseInfosRes[2],
		Status: baseInfosRes[3],
	}

	// StatusOther
	restOfTheLine := lineStr[len(baseInfosRes[0]):]
	if len(restOfTheLine) > 0 {
		statusOtherExp := regexp.MustCompile(`\((?P<statusOther>[a-zA-Z ,]*)\)`)
		statusOtherRes := statusOtherExp.FindStringSubmatch(restOfTheLine)
		if statusOtherRes != nil {
			simInfo.StatusOther = statusOtherRes[1]
		}
	}
	return simInfo, nil
}

func collectAllSimIDs(simctlListOutputToScan string) models.SimulatorsGroupedByIOSVersionsModel {
	simulatorsByIOSVersions := models.SimulatorsGroupedByIOSVersionsModel{}
	currIOSVersion := ""

	fscanner := bufio.NewScanner(strings.NewReader(simctlListOutputToScan))
	isDevicesSectionFound := false
	for fscanner.Scan() {
		aLine := fscanner.Text()

		if aLine == "== Devices ==" {
			isDevicesSectionFound = true
			continue
		}

		if !isDevicesSectionFound {
			continue
		}
		if strings.HasPrefix(aLine, "==") {
			isDevicesSectionFound = false
			continue
		}
		if strings.HasPrefix(aLine, "--") {
			iosVersionSectionExp := regexp.MustCompile(`-- (?P<iosVersionSection>.*) --`)
			iosVersionSectionRes := iosVersionSectionExp.FindStringSubmatch(aLine)
			if iosVersionSectionRes != nil {
				currIOSVersion = iosVersionSectionRes[1]
			}
			continue
		}

		// fmt.Println("-> ", aLine)
		simInfo, err := getSimInfoFromLine(aLine)
		if err != nil {
			fmt.Println(" [!] Error scanning the line for Simulator info: ", err)
		}

		currIOSVersionSimList := simulatorsByIOSVersions[currIOSVersion]
		currIOSVersionSimList = append(currIOSVersionSimList, simInfo)
		simulatorsByIOSVersions[currIOSVersion] = currIOSVersionSimList
	}

	return simulatorsByIOSVersions
}

// Compares sematic versions with 2 components (9.1, 9.2, ...)
// Return true if first version is greater then second
func isOsVersionGreater(osVersion, otherOsVersion string) (bool, error) {
	versionPtrs := []*ver.Version{}
	for _, osVer := range []string{osVersion, otherOsVersion} {
		osVersionSplit := strings.Split(osVer, " ")
		if len(osVersionSplit) != 2 {
			return false, fmt.Errorf("failed to parse version: %s", osVer)
		}

		version, err := ver.NewVersion(osVersionSplit[1])
		if err != nil {
			return false, err
		}

		versionPtrs = append(versionPtrs, version)
	}

	return versionPtrs[0].GreaterThan(versionPtrs[1]), nil
}

func getLatestOsVersion(desiredPlatform, desiredDevice string, allSimIDsGroupedBySimVersion models.SimulatorsGroupedByIOSVersionsModel) (string, error) {
	latestOsVersion := ""
	for osVersion, simInfos := range allSimIDsGroupedBySimVersion {
		if !strings.HasPrefix(osVersion, desiredPlatform) {
			continue
		}

		deviceExist := false
		for _, simInfo := range simInfos {
			if simInfo.Name == desiredDevice {
				deviceExist = true
				break
			}
		}

		if !deviceExist {
			continue
		}

		if latestOsVersion == "" {
			latestOsVersion = osVersion
		} else {
			greater, err := isOsVersionGreater(osVersion, latestOsVersion)
			if err != nil {
				return "", err
			}

			if greater {
				latestOsVersion = osVersion
			}
		}
	}

	if latestOsVersion == "" {
		return "", fmt.Errorf("failed to find latest os version for (%s) - (%s)", desiredPlatform, desiredDevice)
	}

	return latestOsVersion, nil
}

//=======================================
// Main
//=======================================

// GetSimulator ...
func GetSimulator(simulatorPlatform, simulatorDevice, simulatorOsVersion string) (models.SimInfoModel, error) {
	cmd := exec.Command("xcrun", "simctl", "list")
	outBytes, err := cmd.CombinedOutput()
	if err != nil {
		return models.SimInfoModel{}, err
	}

	simctlListOut := string(outBytes)

	allSimIDsGroupedBySimVersion := collectAllSimIDs(simctlListOut)

	//
	// map desired inputs
	simulatorPlatformSplit := strings.Split(simulatorPlatform, " Simulator")
	if len(simulatorPlatformSplit) == 0 {
		return models.SimInfoModel{}, fmt.Errorf("failed to parse simulator platform (%s)", simulatorPlatform)
	}

	if simulatorDevice == "iPad" {
		log.LogWarn("Given device (%s) is deprecated, using (iPad 2)...", simulatorDevice)
		simulatorDevice = "iPad 2"
	}

	desiredPlatform := simulatorPlatformSplit[0]
	desiredOsVersion := ""

	if simulatorOsVersion == "latest" {
		latestOsVersion, err := getLatestOsVersion(desiredPlatform, simulatorDevice, allSimIDsGroupedBySimVersion)
		if err != nil {
			return models.SimInfoModel{}, fmt.Errorf("failed to fer latest os version, error: %s", err)
		}
		desiredOsVersion = latestOsVersion
	} else {
		desiredOsVersion = fmt.Sprintf("%s %s", desiredPlatform, simulatorOsVersion)
	}

	//
	// find desired simulator
	simInfoList, found := allSimIDsGroupedBySimVersion[desiredOsVersion]
	if !found {
		return models.SimInfoModel{}, fmt.Errorf("no simulator found for desired os: %s", desiredOsVersion)
	}

	for _, simInfo := range simInfoList {
		if simInfo.Name == simulatorDevice {
			return simInfo, nil
		}
	}

	return models.SimInfoModel{}, fmt.Errorf("%s - %s - %s not found", simulatorPlatform, simulatorDevice, simulatorOsVersion)
}

// BootSimulator ...
func BootSimulator(simulator models.SimInfoModel, xcodebuildVersion models.XcodebuildVersionModel) error {
	simulatorApp := "Simulator"
	if xcodebuildVersion.MajorVersion == 6 {
		simulatorApp = "iOS Simulator"
	}

	openCmd := exec.Command("open", "-a", simulatorApp, "--args", "-CurrentDeviceUDID", simulator.SimID)

	log.LogDetails("$ %s", cmd.PrintableCommandArgs(openCmd.Args))

	out, err := openCmd.CombinedOutput()
	outStr := string(out)
	if err != nil {
		return fmt.Errorf("failed to start simulators (%s), output: %s, error: %s", simulator.SimID, outStr, err)
	}

	return nil
}

// GetXcodeVersion ...
func GetXcodeVersion() (models.XcodebuildVersionModel, error) {
	cmd := exec.Command("xcodebuild", "-version")
	outBytes, err := cmd.CombinedOutput()
	outStr := string(outBytes)
	if err != nil {
		return models.XcodebuildVersionModel{}, fmt.Errorf("xcodebuild -version failed, err: %s, details: %s", err, outStr)
	}

	split := strings.Split(outStr, "\n")
	if len(split) == 0 {
		return models.XcodebuildVersionModel{}, fmt.Errorf("failed to parse xcodebuild version output (%s)", outStr)
	}

	xcodebuildVersion := split[0]
	buildVersion := split[1]

	split = strings.Split(xcodebuildVersion, " ")
	if len(split) != 2 {
		return models.XcodebuildVersionModel{}, fmt.Errorf("failed to parse xcodebuild version output (%s)", outStr)
	}

	version := split[1]

	split = strings.Split(version, ".")
	majorVersionStr := split[0]

	majorVersion, err := strconv.ParseInt(majorVersionStr, 10, 32)
	if err != nil {
		return models.XcodebuildVersionModel{}, fmt.Errorf("failed to parse xcodebuild version output (%s), error: %s", outStr, err)
	}

	return models.XcodebuildVersionModel{
		Version:      xcodebuildVersion,
		BuildVersion: buildVersion,
		MajorVersion: majorVersion,
	}, nil
}
<commit_msg>BootSimulator - multi xcode bug fix<commit_after>package xcodeutil

import (
	"bufio"
	"fmt"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	cmd "github.com/bitrise-io/xcode-test/command"
	log "github.com/bitrise-io/xcode-test/logutil"
	"github.com/bitrise-io/xcode-test/models"
	ver "github.com/hashicorp/go-version"
)

var (
	osRegexp          = regexp.MustCompile(`-- (.+) (\d\.\d) --`)
	deviceStateRegexp = regexp.MustCompile(` *(.+) \(([A-Z0-9-]+)\) \((.+)\)`)
)

//=======================================
// Utility
//=======================================

// a simulator info line should look like this:
//  iPhone 5s (EA1C7E48-8137-428C-A0A5-B2C63FF276EB) (Shutdown)
// or
//  iPhone 4s (51B10EBD-C949-49F5-A38B-E658F41640FF) (Shutdown) (unavailable, runtime profile not found)
func getSimInfoFromLine(lineStr string) (models.SimInfoModel, error) {
	baseInfosExp := regexp.MustCompile(`(?P<deviceName>[a-zA-Z].*[a-zA-Z0-9 -]*) \((?P<simulatorID>[a-zA-Z0-9-]{36})\) \((?P<status>[a-zA-Z]*)\)`)
	baseInfosRes := baseInfosExp.FindStringSubmatch(lineStr)
	if baseInfosRes == nil {
		return models.SimInfoModel{}, fmt.Errorf("No match found")
	}

	simInfo := models.SimInfoModel{
		Name:   baseInfosRes[1],
		SimID:  baseInfosRes[2],
		Status: baseInfosRes[3],
	}

	// StatusOther
	restOfTheLine := lineStr[len(baseInfosRes[0]):]
	if len(restOfTheLine) > 0 {
		statusOtherExp := regexp.MustCompile(`\((?P<statusOther>[a-zA-Z ,]*)\)`)
		statusOtherRes := statusOtherExp.FindStringSubmatch(restOfTheLine)
		if statusOtherRes != nil {
			simInfo.StatusOther = statusOtherRes[1]
		}
	}
	return simInfo, nil
}

func collectAllSimIDs(simctlListOutputToScan string) models.SimulatorsGroupedByIOSVersionsModel {
	simulatorsByIOSVersions := models.SimulatorsGroupedByIOSVersionsModel{}
	currIOSVersion := ""

	fscanner := bufio.NewScanner(strings.NewReader(simctlListOutputToScan))
	isDevicesSectionFound := false
	for fscanner.Scan() {
		aLine := fscanner.Text()

		if aLine == "== Devices ==" {
			isDevicesSectionFound = true
			continue
		}

		if !isDevicesSectionFound {
			continue
		}
		if strings.HasPrefix(aLine, "==") {
			isDevicesSectionFound = false
			continue
		}
		if strings.HasPrefix(aLine, "--") {
			iosVersionSectionExp := regexp.MustCompile(`-- (?P<iosVersionSection>.*) --`)
			iosVersionSectionRes := iosVersionSectionExp.FindStringSubmatch(aLine)
			if iosVersionSectionRes != nil {
				currIOSVersion = iosVersionSectionRes[1]
			}
			continue
		}

		// fmt.Println("-> ", aLine)
		simInfo, err := getSimInfoFromLine(aLine)
		if err != nil {
			fmt.Println(" [!] Error scanning the line for Simulator info: ", err)
		}

		currIOSVersionSimList := simulatorsByIOSVersions[currIOSVersion]
		currIOSVersionSimList = append(currIOSVersionSimList, simInfo)
		simulatorsByIOSVersions[currIOSVersion] = currIOSVersionSimList
	}

	return simulatorsByIOSVersions
}

// Compares sematic versions with 2 components (9.1, 9.2, ...)
// Return true if first version is greater then second
func isOsVersionGreater(osVersion, otherOsVersion string) (bool, error) {
	versionPtrs := []*ver.Version{}
	for _, osVer := range []string{osVersion, otherOsVersion} {
		osVersionSplit := strings.Split(osVer, " ")
		if len(osVersionSplit) != 2 {
			return false, fmt.Errorf("failed to parse version: %s", osVer)
		}

		version, err := ver.NewVersion(osVersionSplit[1])
		if err != nil {
			return false, err
		}

		versionPtrs = append(versionPtrs, version)
	}

	return versionPtrs[0].GreaterThan(versionPtrs[1]), nil
}

func getLatestOsVersion(desiredPlatform, desiredDevice string, allSimIDsGroupedBySimVersion models.SimulatorsGroupedByIOSVersionsModel) (string, error) {
	latestOsVersion := ""
	for osVersion, simInfos := range allSimIDsGroupedBySimVersion {
		if !strings.HasPrefix(osVersion, desiredPlatform) {
			continue
		}

		deviceExist := false
		for _, simInfo := range simInfos {
			if simInfo.Name == desiredDevice {
				deviceExist = true
				break
			}
		}

		if !deviceExist {
			continue
		}

		if latestOsVersion == "" {
			latestOsVersion = osVersion
		} else {
			greater, err := isOsVersionGreater(osVersion, latestOsVersion)
			if err != nil {
				return "", err
			}

			if greater {
				latestOsVersion = osVersion
			}
		}
	}

	if latestOsVersion == "" {
		return "", fmt.Errorf("failed to find latest os version for (%s) - (%s)", desiredPlatform, desiredDevice)
	}

	return latestOsVersion, nil
}

//=======================================
// Main
//=======================================

// GetSimulator ...
func GetSimulator(simulatorPlatform, simulatorDevice, simulatorOsVersion string) (models.SimInfoModel, error) {
	cmd := exec.Command("xcrun", "simctl", "list")
	outBytes, err := cmd.CombinedOutput()
	if err != nil {
		return models.SimInfoModel{}, err
	}

	simctlListOut := string(outBytes)

	allSimIDsGroupedBySimVersion := collectAllSimIDs(simctlListOut)

	//
	// map desired inputs
	simulatorPlatformSplit := strings.Split(simulatorPlatform, " Simulator")
	if len(simulatorPlatformSplit) == 0 {
		return models.SimInfoModel{}, fmt.Errorf("failed to parse simulator platform (%s)", simulatorPlatform)
	}

	if simulatorDevice == "iPad" {
		log.LogWarn("Given device (%s) is deprecated, using (iPad 2)...", simulatorDevice)
		simulatorDevice = "iPad 2"
	}

	desiredPlatform := simulatorPlatformSplit[0]
	desiredOsVersion := ""

	if simulatorOsVersion == "latest" {
		latestOsVersion, err := getLatestOsVersion(desiredPlatform, simulatorDevice, allSimIDsGroupedBySimVersion)
		if err != nil {
			return models.SimInfoModel{}, fmt.Errorf("failed to fer latest os version, error: %s", err)
		}
		desiredOsVersion = latestOsVersion
	} else {
		desiredOsVersion = fmt.Sprintf("%s %s", desiredPlatform, simulatorOsVersion)
	}

	//
	// find desired simulator
	simInfoList, found := allSimIDsGroupedBySimVersion[desiredOsVersion]
	if !found {
		return models.SimInfoModel{}, fmt.Errorf("no simulator found for desired os: %s", desiredOsVersion)
	}

	for _, simInfo := range simInfoList {
		if simInfo.Name == simulatorDevice {
			return simInfo, nil
		}
	}

	return models.SimInfoModel{}, fmt.Errorf("%s - %s - %s not found", simulatorPlatform, simulatorDevice, simulatorOsVersion)
}

func getXcodeDeveloperDirPath() (string, error) {
	cmd := exec.Command("xcode-select", "--print-path")
	outBytes, err := cmd.CombinedOutput()
	outStr := string(outBytes)
	return strings.TrimSpace(outStr), err
}

// BootSimulator ...
func BootSimulator(simulator models.SimInfoModel, xcodebuildVersion models.XcodebuildVersionModel) error {
	simulatorApp := "Simulator"
	if xcodebuildVersion.MajorVersion == 6 {
		simulatorApp = "iOS Simulator"
	}
	xcodeDevDirPth, err := getXcodeDeveloperDirPath()
	if err != nil {
		return fmt.Errorf("Failed to get Xcode Developer Directory - most likely Xcode.app is not installed")
	}
	simulatorAppFullPath := filepath.Join(xcodeDevDirPth, "Applications", simulatorApp+".app")

	openCmd := exec.Command("open", simulatorAppFullPath, "--args", "-CurrentDeviceUDID", simulator.SimID)

	log.LogDetails("$ %s", cmd.PrintableCommandArgs(openCmd.Args))

	out, err := openCmd.CombinedOutput()
	outStr := string(out)
	if err != nil {
		return fmt.Errorf("failed to start simulators (%s), output: %s, error: %s", simulator.SimID, outStr, err)
	}

	return nil
}

// GetXcodeVersion ...
func GetXcodeVersion() (models.XcodebuildVersionModel, error) {
	cmd := exec.Command("xcodebuild", "-version")
	outBytes, err := cmd.CombinedOutput()
	outStr := string(outBytes)
	if err != nil {
		return models.XcodebuildVersionModel{}, fmt.Errorf("xcodebuild -version failed, err: %s, details: %s", err, outStr)
	}

	split := strings.Split(outStr, "\n")
	if len(split) == 0 {
		return models.XcodebuildVersionModel{}, fmt.Errorf("failed to parse xcodebuild version output (%s)", outStr)
	}

	xcodebuildVersion := split[0]
	buildVersion := split[1]

	split = strings.Split(xcodebuildVersion, " ")
	if len(split) != 2 {
		return models.XcodebuildVersionModel{}, fmt.Errorf("failed to parse xcodebuild version output (%s)", outStr)
	}

	version := split[1]

	split = strings.Split(version, ".")
	majorVersionStr := split[0]

	majorVersion, err := strconv.ParseInt(majorVersionStr, 10, 32)
	if err != nil {
		return models.XcodebuildVersionModel{}, fmt.Errorf("failed to parse xcodebuild version output (%s), error: %s", outStr, err)
	}

	return models.XcodebuildVersionModel{
		Version:      xcodebuildVersion,
		BuildVersion: buildVersion,
		MajorVersion: majorVersion,
	}, nil
}
<|endoftext|>


<commit_before>// This package is created for deleting the old guest accounts
package main

import (
	"flag"
	"koding/db/models"
	helper "koding/db/mongodb/modelhelper"
	"koding/helpers"
	"koding/tools/config"
	"koding/tools/logger"
	"time"

	"labix.org/v2/mgo/bson"
)

var log = logger.New("Guest cleaner worker")

var (
	configProfile = flag.String("c", "", "Configuration profile from file")
	flagSkip      = flag.Int("s", 0, "Configuration profile from file")
	flagLimit     = flag.Int("l", 1000, "Configuration profile from file")
)

func initialize() {
	flag.Parse()
	if *configProfile == "" {
		log.Fatal("Please define config file with -c")
	}

	conf := config.MustConfig(*configProfile)
	helper.Initialize(conf.Mongo)
}

func main() {
	// init the package
	initialize()

	log.Info("Guest cleaner worker started")

	iterOptions := helpers.NewIterOptions()
	iterOptions.CollectionName = "jAccounts"
	iterOptions.F = deleteGuestAccounts
	iterOptions.Filter = createFilter()
	iterOptions.DataType = &models.Account{}
	iterOptions.Limit = *flagLimit
	iterOptions.Skip = *flagSkip
	iterOptions.Log = log

	log.SetLevel(logger.DEBUG)

	err := helpers.Iter(helper.Mongo, iterOptions)
	if err != nil {
		log.Fatal("Error while iter: %v", err)
	}
	log.Info("Guest cleaner worker finished")
}

func createFilter() helper.Selector {
	oneHourAgo := time.Now().Add(-time.Minute * 60).UTC()
	return helper.Selector{
		"type":           "unregistered",
		"status":         helper.Selector{"$ne": "tobedeleted"},
		"meta.createdAt": helper.Selector{"$lte": oneHourAgo},
	}
}

// deleteGuestAccounts deletes all related documents and
// relationships for given account
func deleteGuestAccounts(account interface{}) error {
	acc := account.(*models.Account)
	log.Info("Deleting account: %v", acc.Profile.Nickname)

	// clear all sessions for acc
	clearAllSessions(acc)

	// clear all JNames for current acc
	clearAllNames(acc)

	selector := helper.Selector{"$or": []helper.Selector{
		helper.Selector{"sourceId": acc.Id},
		helper.Selector{"targetId": acc.Id},
	}}

	// Get all relationships for current acc
	rels, err := helper.GetAllRelationships(selector)
	if err != nil {
		log.Error("Error while getting the all relationships for deletion %v", err)
		return nil
	}

	// Delete documents with all their relationships
	deleteDocumentsWithRelationships(rels)

	// now time to delete acc itself
	if err := helper.RemoveAccount(acc.Id); err != nil {
		log.Error("Error while deleting the acc %v", err)
		return nil
	}
}

// clearAllSessions deletes related JSesssions from database
func clearAllSessions(account *models.Account) {
	selector := helper.Selector{"username": account.Profile.Nickname}
	if err := helper.RemoveAllDocuments("jSessions", selector); err != nil {
		log.Error("Error while deleting from JSession %v", err)
	}
}

// clearAllNames deletes related documents from JName collection
func clearAllNames(account *models.Account) {
	selector := helper.Selector{"name": account.Profile.Nickname}
	if err := helper.RemoveAllDocuments("jName", selector); err != nil {
		log.Error("Error while deleting from JName %v", err)
	}
}

// deleteDocumentsWithRelationships accepts a slice of Relationship
// and range over it, for each relationship checks if source or the target
// is eligible for deletion, at the end deletes the relationship itself
func deleteDocumentsWithRelationships(rels []models.Relationship) {
	if len(rels) == 0 {
		log.Debug("no item to process")
		return
	}

	for _, rel := range rels {
		// delete source first
		if checkIfEligibleToDelete(rel.SourceName) {
			deleteDocument(helper.GetCollectionName(rel.SourceName), rel.SourceId)
		}
		// then delete target
		if checkIfEligibleToDelete(rel.TargetName) {
			deleteDocument(helper.GetCollectionName(rel.TargetName), rel.TargetId)
		}
		// delete relationship itself
		deleteDocument("relationships", rel.Id)
	}
}

// safe to delete these documents automatically
var whiteListedModels = []string{"JSession", "JUser", "JVM", "JDomain", "JAppStorage", "JName"}

//checkIfEligibleToDelete ranges over whiteListedModels and returns bool as checking
// name is in that list or not
func checkIfEligibleToDelete(modelName string) bool {
	for _, name := range whiteListedModels {
		if name == modelName {
			return true
		}
	}

	log.Debug("not eligible %v", modelName)
	return false
}

// deleteDocument deletes the given id from given collection
// this is here just for logging purposes
func deleteDocument(collectionName string, id bson.ObjectId) {
	log.Debug("removing id: %v from collectionName: %v ", id.Hex(), collectionName)
	if err := helper.RemoveDocument(collectionName, id); err != nil {
		log.Error(
			"couldnt remove collectionId: %v from collectionName: %v  Error: %v",
			id.Hex(),
			collectionName,
			err,
		)
	}
}
<commit_msg>GuestCleaner: add forgotten error<commit_after>// This package is created for deleting the old guest accounts
package main

import (
	"flag"
	"koding/db/models"
	helper "koding/db/mongodb/modelhelper"
	"koding/helpers"
	"koding/tools/config"
	"koding/tools/logger"
	"time"

	"labix.org/v2/mgo/bson"
)

var log = logger.New("Guest cleaner worker")

var (
	configProfile = flag.String("c", "", "Configuration profile from file")
	flagSkip      = flag.Int("s", 0, "Configuration profile from file")
	flagLimit     = flag.Int("l", 1000, "Configuration profile from file")
)

func initialize() {
	flag.Parse()
	if *configProfile == "" {
		log.Fatal("Please define config file with -c")
	}

	conf := config.MustConfig(*configProfile)
	helper.Initialize(conf.Mongo)
}

func main() {
	// init the package
	initialize()

	log.Info("Guest cleaner worker started")

	iterOptions := helpers.NewIterOptions()
	iterOptions.CollectionName = "jAccounts"
	iterOptions.F = deleteGuestAccounts
	iterOptions.Filter = createFilter()
	iterOptions.DataType = &models.Account{}
	iterOptions.Limit = *flagLimit
	iterOptions.Skip = *flagSkip
	iterOptions.Log = log

	log.SetLevel(logger.DEBUG)

	err := helpers.Iter(helper.Mongo, iterOptions)
	if err != nil {
		log.Fatal("Error while iter: %v", err)
	}
	log.Info("Guest cleaner worker finished")
}

func createFilter() helper.Selector {
	oneHourAgo := time.Now().Add(-time.Minute * 60).UTC()
	return helper.Selector{
		"type":           "unregistered",
		"status":         helper.Selector{"$ne": "tobedeleted"},
		"meta.createdAt": helper.Selector{"$lte": oneHourAgo},
	}
}

// deleteGuestAccounts deletes all related documents and
// relationships for given account
func deleteGuestAccounts(account interface{}) error {
	acc := account.(*models.Account)
	log.Info("Deleting account: %v", acc.Profile.Nickname)

	// clear all sessions for acc
	clearAllSessions(acc)

	// clear all JNames for current acc
	clearAllNames(acc)

	selector := helper.Selector{"$or": []helper.Selector{
		helper.Selector{"sourceId": acc.Id},
		helper.Selector{"targetId": acc.Id},
	}}

	// Get all relationships for current acc
	rels, err := helper.GetAllRelationships(selector)
	if err != nil {
		log.Error("Error while getting the all relationships for deletion %v", err)
		return nil
	}

	// Delete documents with all their relationships
	deleteDocumentsWithRelationships(rels)

	// now time to delete acc itself
	if err := helper.RemoveAccount(acc.Id); err != nil {
		log.Error("Error while deleting the acc %v", err)
		return nil
	}
	return nil
}

// clearAllSessions deletes related JSesssions from database
func clearAllSessions(account *models.Account) {
	selector := helper.Selector{"username": account.Profile.Nickname}
	if err := helper.RemoveAllDocuments("jSessions", selector); err != nil {
		log.Error("Error while deleting from JSession %v", err)
	}
}

// clearAllNames deletes related documents from JName collection
func clearAllNames(account *models.Account) {
	selector := helper.Selector{"name": account.Profile.Nickname}
	if err := helper.RemoveAllDocuments("jName", selector); err != nil {
		log.Error("Error while deleting from JName %v", err)
	}
}

// deleteDocumentsWithRelationships accepts a slice of Relationship
// and range over it, for each relationship checks if source or the target
// is eligible for deletion, at the end deletes the relationship itself
func deleteDocumentsWithRelationships(rels []models.Relationship) {
	if len(rels) == 0 {
		log.Debug("no item to process")
		return
	}

	for _, rel := range rels {
		// delete source first
		if checkIfEligibleToDelete(rel.SourceName) {
			deleteDocument(helper.GetCollectionName(rel.SourceName), rel.SourceId)
		}
		// then delete target
		if checkIfEligibleToDelete(rel.TargetName) {
			deleteDocument(helper.GetCollectionName(rel.TargetName), rel.TargetId)
		}
		// delete relationship itself
		deleteDocument("relationships", rel.Id)
	}
}

// safe to delete these documents automatically
var whiteListedModels = []string{"JSession", "JUser", "JVM", "JDomain", "JAppStorage", "JName"}

//checkIfEligibleToDelete ranges over whiteListedModels and returns bool as checking
// name is in that list or not
func checkIfEligibleToDelete(modelName string) bool {
	for _, name := range whiteListedModels {
		if name == modelName {
			return true
		}
	}

	log.Debug("not eligible %v", modelName)
	return false
}

// deleteDocument deletes the given id from given collection
// this is here just for logging purposes
func deleteDocument(collectionName string, id bson.ObjectId) {
	log.Debug("removing id: %v from collectionName: %v ", id.Hex(), collectionName)
	if err := helper.RemoveDocument(collectionName, id); err != nil {
		log.Error(
			"couldnt remove collectionId: %v from collectionName: %v  Error: %v",
			id.Hex(),
			collectionName,
			err,
		)
	}
}
<|endoftext|>


<commit_before>package cluster

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	// Force-refresh the state of the node this often.
	stateRefreshPeriod = 30 * time.Second

	// Timeout for requests sent out to the node.
	requestTimeout = 10 * time.Second
)

func NewNode(addr string, overcommitRatio float64) *Node {
	e := &Node{
		Addr:            addr,
		Labels:          make(map[string]string),
		ch:              make(chan bool),
		containers:      make(map[string]*Container),
		healthy:         true,
		overcommitRatio: int64(overcommitRatio * 100),
	}
	return e
}

type Node struct {
	sync.RWMutex

	ID     string
	IP     string
	Addr   string
	Name   string
	Cpus   int64
	Memory int64
	Labels map[string]string

	ch              chan bool
	containers      map[string]*Container
	images          []*dockerclient.Image
	client          dockerclient.Client
	eventHandler    EventHandler
	healthy         bool
	overcommitRatio int64
}

// Connect will initialize a connection to the Docker daemon running on the
// host, gather machine specs (memory, cpu, ...) and monitor state changes.
func (n *Node) Connect(config *tls.Config) error {
	host, port, err := net.SplitHostPort(n.Addr)
	if err != nil {
		return err
	}

	addr, err := net.ResolveIPAddr("ip4", host)
	if err != nil {
		return err
	}
	n.IP = addr.IP.String()

	c, err := dockerclient.NewDockerClientTimeout(n.IP+":"+port, config, time.Duration(requestTimeout))
	if err != nil {
		return err
	}

	return n.connectClient(c)
}

func (n *Node) connectClient(client dockerclient.Client) error {
	n.client = client

	// Fetch the engine labels.
	if err := n.updateSpecs(); err != nil {
		n.client = nil
		return err
	}

	// Force a state update before returning.
	if err := n.RefreshContainers(true); err != nil {
		n.client = nil
		return err
	}

	if err := n.refreshImages(); err != nil {
		n.client = nil
		return err
	}

	// Start the update loop.
	go n.refreshLoop()

	// Start monitoring events from the Node.
	n.client.StartMonitorEvents(n.handler, nil)
	n.emitEvent("node_connect")

	return nil
}

// IsConnected returns true if the engine is connected to a remote docker API
func (n *Node) IsConnected() bool {
	return n.client != nil
}

func (n *Node) IsHealthy() bool {
	return n.healthy
}

// Gather node specs (CPU, memory, constraints, ...).
func (n *Node) updateSpecs() error {
	info, err := n.client.Info()
	if err != nil {
		return err
	}
	// Older versions of Docker don't expose the ID field and are not supported
	// by Swarm.  Catch the error ASAP and refuse to connect.
	if len(info.ID) == 0 {
		return fmt.Errorf("Node %s is running an unsupported version of Docker Engine. Please upgrade.", n.Addr)
	}
	n.ID = info.ID
	n.Name = info.Name
	n.Cpus = info.NCPU
	n.Memory = info.MemTotal
	n.Labels = map[string]string{
		"storagedriver":   info.Driver,
		"executiondriver": info.ExecutionDriver,
		"kernelversion":   info.KernelVersion,
		"operatingsystem": info.OperatingSystem,
	}
	for _, label := range info.Labels {
		kv := strings.SplitN(label, "=", 2)
		n.Labels[kv[0]] = kv[1]
	}
	return nil
}

// Refresh the list of images on the node.
func (n *Node) refreshImages() error {
	images, err := n.client.ListImages()
	if err != nil {
		return err
	}
	n.Lock()
	n.images = images
	n.Unlock()
	return nil
}

// Refresh the list and status of containers running on the node. If `full` is
// true, each container will be inspected.
func (n *Node) RefreshContainers(full bool) error {
	containers, err := n.client.ListContainers(true, false, "")
	if err != nil {
		return err
	}

	merged := make(map[string]*Container)
	for _, c := range containers {
		merged, err = n.updateContainer(c, merged, full)
		if err != nil {
			log.WithFields(log.Fields{"name": n.Name, "id": n.ID}).Errorf("Unable to update state of container %q", c.Id)
		}
	}

	n.Lock()
	defer n.Unlock()
	n.containers = merged

	log.WithFields(log.Fields{"id": n.ID, "name": n.Name}).Debugf("Updated node state")
	return nil
}

// Refresh the status of a container running on the node. If `full` is true,
// the container will be inspected.
func (n *Node) RefreshContainer(ID string, full bool) error {
	containers, err := n.client.ListContainers(true, false, fmt.Sprintf("{%q:[%q]}", "id", ID))
	if err != nil {
		return err
	}

	if len(containers) > 1 {
		// We expect one container, if we get more than one, trigger a full refresh.
		return n.RefreshContainers(full)
	}

	if len(containers) == 0 {
		// The container doesn't exist on the node, remove it.
		n.Lock()
		delete(n.containers, ID)
		n.Unlock()

		return nil
	}

	_, err = n.updateContainer(containers[0], n.containers, full)
	return err
}

func (n *Node) updateContainer(c dockerclient.Container, containers map[string]*Container, full bool) (map[string]*Container, error) {
	var container *Container

	n.Lock()

	if current, exists := n.containers[c.Id]; exists {
		// The container is already known.
		container = current
	} else {
		// This is a brand new container. We need to do a full refresh.
		container = &Container{
			Node: n,
		}
		full = true
	}

	// Update its internal state.
	container.Container = c
	containers[container.Id] = container

	// Release the lock here as the next step is slow.
	n.Unlock()

	// Update ContainerInfo.
	if full {
		info, err := n.client.InspectContainer(c.Id)
		if err != nil {
			return nil, err
		}
		container.Info = *info
		// real CpuShares -> nb of CPUs
		container.Info.Config.CpuShares = container.Info.Config.CpuShares / 100.0 * n.Cpus
	}

	return containers, nil
}

func (n *Node) RefreshContainersAsync() {
	n.ch <- true
}

func (n *Node) refreshLoop() {
	for {
		var err error
		select {
		case <-n.ch:
			err = n.RefreshContainers(false)
		case <-time.After(stateRefreshPeriod):
			err = n.RefreshContainers(false)
		}

		if err == nil {
			err = n.refreshImages()
		}

		if err != nil {
			if n.healthy {
				n.emitEvent("node_disconnect")
			}
			n.healthy = false
			log.WithFields(log.Fields{"name": n.Name, "id": n.ID}).Errorf("Flagging node as dead. Updated state failed: %v", err)
		} else {
			if !n.healthy {
				log.WithFields(log.Fields{"name": n.Name, "id": n.ID}).Info("Node came back to life. Hooray!")
				n.client.StopAllMonitorEvents()
				n.client.StartMonitorEvents(n.handler, nil)
				n.emitEvent("node_reconnect")
				if err := n.updateSpecs(); err != nil {
					log.WithFields(log.Fields{"name": n.Name, "id": n.ID}).Errorf("Update node specs failed: %v", err)
				}
			}
			n.healthy = true
		}
	}
}

func (n *Node) emitEvent(event string) {
	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}
	ev := &Event{
		Event: dockerclient.Event{
			Status: event,
			From:   "swarm",
			Time:   time.Now().Unix(),
		},
		Node: n,
	}
	n.eventHandler.Handle(ev)
}

// Return the sum of memory reserved by containers.
func (n *Node) ReservedMemory() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.Memory
	}
	n.RUnlock()
	return r
}

// Return the sum of CPUs reserved by containers.
func (n *Node) ReservedCpus() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.CpuShares
	}
	n.RUnlock()
	return r
}

func (n *Node) UsableMemory() int64 {
	return n.Memory + (n.Memory * n.overcommitRatio / 100)
}

func (n *Node) UsableCpus() int64 {
	return n.Cpus + (n.Cpus * n.overcommitRatio / 100)
}

func (n *Node) Create(config *dockerclient.ContainerConfig, name string, pullImage bool) (*Container, error) {
	var (
		err    error
		id     string
		client = n.client
	)

	newConfig := *config

	// nb of CPUs -> real CpuShares
	newConfig.CpuShares = config.CpuShares * 100 / n.Cpus

	if id, err = client.CreateContainer(&newConfig, name); err != nil {
		// If the error is other than not found, abort immediately.
		if err != dockerclient.ErrNotFound || !pullImage {
			return nil, err
		}
		// Otherwise, try to pull the image...
		if err = n.Pull(config.Image); err != nil {
			return nil, err
		}
		// ...And try again.
		if id, err = client.CreateContainer(&newConfig, name); err != nil {
			return nil, err
		}
	}

	// Register the container immediately while waiting for a state refresh.
	// Force a state refresh to pick up the newly created container.
	n.RefreshContainer(id, true)

	n.RLock()
	defer n.RUnlock()

	return n.containers[id], nil
}

// Destroy and remove a container from the node.
func (n *Node) Destroy(container *Container, force bool) error {
	if err := n.client.RemoveContainer(container.Id, force); err != nil {
		return err
	}

	// Remove the container from the state. Eventually, the state refresh loop
	// will rewrite this.
	n.Lock()
	defer n.Unlock()
	delete(n.containers, container.Id)

	return nil
}

func (n *Node) Pull(image string) error {
	if err := n.client.PullImage(image, nil); err != nil {
		return err
	}
	return nil
}

// Register an event handler.
func (n *Node) Events(h EventHandler) error {
	if n.eventHandler != nil {
		return errors.New("event handler already set")
	}
	n.eventHandler = h
	return nil
}

func (n *Node) Containers() []*Container {
	containers := []*Container{}
	n.RLock()
	for _, container := range n.containers {
		containers = append(containers, container)
	}
	n.RUnlock()
	return containers
}

func (n *Node) Images() []*dockerclient.Image {
	images := []*dockerclient.Image{}
	n.RLock()
	for _, image := range n.images {
		images = append(images, image)
	}
	n.RUnlock()
	return images
}

// Image returns the image with IdOrName in the node
func (n *Node) Image(IdOrName string) *dockerclient.Image {
	size := len(IdOrName)
	for _, image := range n.Images() {
		if image.Id == IdOrName || (size > 2 && strings.HasPrefix(image.Id, IdOrName)) {
			return image
		}
		for _, t := range image.RepoTags {
			if t == IdOrName || (size > 2 && strings.HasPrefix(t, IdOrName)) {
				return image
			}
		}
	}
	return nil
}

func (n *Node) String() string {
	return fmt.Sprintf("node %s addr %s", n.ID, n.Addr)
}

func (n *Node) handler(ev *dockerclient.Event, _ chan error, args ...interface{}) {
	// Something changed - refresh our internal state.
	switch ev.Status {
	case "pull", "untag", "delete":
		// These events refer to images so there's no need to update
		// containers.
		n.refreshImages()
	case "start", "die":
		// If the container is started or stopped, we have to do an inspect in
		// order to get the new NetworkSettings.
		n.RefreshContainer(ev.Id, true)
	default:
		// Otherwise, do a "soft" refresh of the container.
		n.RefreshContainer(ev.Id, false)
	}

	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}

	event := &Event{
		Node:  n,
		Event: *ev,
	}

	n.eventHandler.Handle(event)
}

// Inject a container into the internal state.
func (n *Node) AddContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; ok {
		return errors.New("container already exists")
	}
	n.containers[container.Id] = container
	return nil
}

// Inject an image into the internal state.
func (n *Node) AddImage(image *dockerclient.Image) {
	n.Lock()
	defer n.Unlock()

	n.images = append(n.images, image)
}

// Remove a container from the internal test.
func (n *Node) RemoveContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; !ok {
		return errors.New("container not found")
	}
	delete(n.containers, container.Id)
	return nil
}

// Wipes the internal container state.
func (n *Node) CleanupContainers() {
	n.Lock()
	n.containers = make(map[string]*Container)
	n.Unlock()
}
<commit_msg>Node: Connect with the provided address instead of resolving the IP.<commit_after>package cluster

import (
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	// Force-refresh the state of the node this often.
	stateRefreshPeriod = 30 * time.Second

	// Timeout for requests sent out to the node.
	requestTimeout = 10 * time.Second
)

func NewNode(addr string, overcommitRatio float64) *Node {
	e := &Node{
		Addr:            addr,
		Labels:          make(map[string]string),
		ch:              make(chan bool),
		containers:      make(map[string]*Container),
		healthy:         true,
		overcommitRatio: int64(overcommitRatio * 100),
	}
	return e
}

type Node struct {
	sync.RWMutex

	ID     string
	IP     string
	Addr   string
	Name   string
	Cpus   int64
	Memory int64
	Labels map[string]string

	ch              chan bool
	containers      map[string]*Container
	images          []*dockerclient.Image
	client          dockerclient.Client
	eventHandler    EventHandler
	healthy         bool
	overcommitRatio int64
}

// Connect will initialize a connection to the Docker daemon running on the
// host, gather machine specs (memory, cpu, ...) and monitor state changes.
func (n *Node) Connect(config *tls.Config) error {
	host, _, err := net.SplitHostPort(n.Addr)
	if err != nil {
		return err
	}

	addr, err := net.ResolveIPAddr("ip4", host)
	if err != nil {
		return err
	}
	n.IP = addr.IP.String()

	c, err := dockerclient.NewDockerClientTimeout("tcp://"+n.Addr, config, time.Duration(requestTimeout))
	if err != nil {
		return err
	}

	return n.connectClient(c)
}

func (n *Node) connectClient(client dockerclient.Client) error {
	n.client = client

	// Fetch the engine labels.
	if err := n.updateSpecs(); err != nil {
		n.client = nil
		return err
	}

	// Force a state update before returning.
	if err := n.RefreshContainers(true); err != nil {
		n.client = nil
		return err
	}

	if err := n.refreshImages(); err != nil {
		n.client = nil
		return err
	}

	// Start the update loop.
	go n.refreshLoop()

	// Start monitoring events from the Node.
	n.client.StartMonitorEvents(n.handler, nil)
	n.emitEvent("node_connect")

	return nil
}

// IsConnected returns true if the engine is connected to a remote docker API
func (n *Node) IsConnected() bool {
	return n.client != nil
}

func (n *Node) IsHealthy() bool {
	return n.healthy
}

// Gather node specs (CPU, memory, constraints, ...).
func (n *Node) updateSpecs() error {
	info, err := n.client.Info()
	if err != nil {
		return err
	}
	// Older versions of Docker don't expose the ID field and are not supported
	// by Swarm.  Catch the error ASAP and refuse to connect.
	if len(info.ID) == 0 {
		return fmt.Errorf("Node %s is running an unsupported version of Docker Engine. Please upgrade.", n.Addr)
	}
	n.ID = info.ID
	n.Name = info.Name
	n.Cpus = info.NCPU
	n.Memory = info.MemTotal
	n.Labels = map[string]string{
		"storagedriver":   info.Driver,
		"executiondriver": info.ExecutionDriver,
		"kernelversion":   info.KernelVersion,
		"operatingsystem": info.OperatingSystem,
	}
	for _, label := range info.Labels {
		kv := strings.SplitN(label, "=", 2)
		n.Labels[kv[0]] = kv[1]
	}
	return nil
}

// Refresh the list of images on the node.
func (n *Node) refreshImages() error {
	images, err := n.client.ListImages()
	if err != nil {
		return err
	}
	n.Lock()
	n.images = images
	n.Unlock()
	return nil
}

// Refresh the list and status of containers running on the node. If `full` is
// true, each container will be inspected.
func (n *Node) RefreshContainers(full bool) error {
	containers, err := n.client.ListContainers(true, false, "")
	if err != nil {
		return err
	}

	merged := make(map[string]*Container)
	for _, c := range containers {
		merged, err = n.updateContainer(c, merged, full)
		if err != nil {
			log.WithFields(log.Fields{"name": n.Name, "id": n.ID}).Errorf("Unable to update state of container %q", c.Id)
		}
	}

	n.Lock()
	defer n.Unlock()
	n.containers = merged

	log.WithFields(log.Fields{"id": n.ID, "name": n.Name}).Debugf("Updated node state")
	return nil
}

// Refresh the status of a container running on the node. If `full` is true,
// the container will be inspected.
func (n *Node) RefreshContainer(ID string, full bool) error {
	containers, err := n.client.ListContainers(true, false, fmt.Sprintf("{%q:[%q]}", "id", ID))
	if err != nil {
		return err
	}

	if len(containers) > 1 {
		// We expect one container, if we get more than one, trigger a full refresh.
		return n.RefreshContainers(full)
	}

	if len(containers) == 0 {
		// The container doesn't exist on the node, remove it.
		n.Lock()
		delete(n.containers, ID)
		n.Unlock()

		return nil
	}

	_, err = n.updateContainer(containers[0], n.containers, full)
	return err
}

func (n *Node) updateContainer(c dockerclient.Container, containers map[string]*Container, full bool) (map[string]*Container, error) {
	var container *Container

	n.Lock()

	if current, exists := n.containers[c.Id]; exists {
		// The container is already known.
		container = current
	} else {
		// This is a brand new container. We need to do a full refresh.
		container = &Container{
			Node: n,
		}
		full = true
	}

	// Update its internal state.
	container.Container = c
	containers[container.Id] = container

	// Release the lock here as the next step is slow.
	n.Unlock()

	// Update ContainerInfo.
	if full {
		info, err := n.client.InspectContainer(c.Id)
		if err != nil {
			return nil, err
		}
		container.Info = *info
		// real CpuShares -> nb of CPUs
		container.Info.Config.CpuShares = container.Info.Config.CpuShares / 100.0 * n.Cpus
	}

	return containers, nil
}

func (n *Node) RefreshContainersAsync() {
	n.ch <- true
}

func (n *Node) refreshLoop() {
	for {
		var err error
		select {
		case <-n.ch:
			err = n.RefreshContainers(false)
		case <-time.After(stateRefreshPeriod):
			err = n.RefreshContainers(false)
		}

		if err == nil {
			err = n.refreshImages()
		}

		if err != nil {
			if n.healthy {
				n.emitEvent("node_disconnect")
			}
			n.healthy = false
			log.WithFields(log.Fields{"name": n.Name, "id": n.ID}).Errorf("Flagging node as dead. Updated state failed: %v", err)
		} else {
			if !n.healthy {
				log.WithFields(log.Fields{"name": n.Name, "id": n.ID}).Info("Node came back to life. Hooray!")
				n.client.StopAllMonitorEvents()
				n.client.StartMonitorEvents(n.handler, nil)
				n.emitEvent("node_reconnect")
				if err := n.updateSpecs(); err != nil {
					log.WithFields(log.Fields{"name": n.Name, "id": n.ID}).Errorf("Update node specs failed: %v", err)
				}
			}
			n.healthy = true
		}
	}
}

func (n *Node) emitEvent(event string) {
	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}
	ev := &Event{
		Event: dockerclient.Event{
			Status: event,
			From:   "swarm",
			Time:   time.Now().Unix(),
		},
		Node: n,
	}
	n.eventHandler.Handle(ev)
}

// Return the sum of memory reserved by containers.
func (n *Node) ReservedMemory() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.Memory
	}
	n.RUnlock()
	return r
}

// Return the sum of CPUs reserved by containers.
func (n *Node) ReservedCpus() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.CpuShares
	}
	n.RUnlock()
	return r
}

func (n *Node) UsableMemory() int64 {
	return n.Memory + (n.Memory * n.overcommitRatio / 100)
}

func (n *Node) UsableCpus() int64 {
	return n.Cpus + (n.Cpus * n.overcommitRatio / 100)
}

func (n *Node) Create(config *dockerclient.ContainerConfig, name string, pullImage bool) (*Container, error) {
	var (
		err    error
		id     string
		client = n.client
	)

	newConfig := *config

	// nb of CPUs -> real CpuShares
	newConfig.CpuShares = config.CpuShares * 100 / n.Cpus

	if id, err = client.CreateContainer(&newConfig, name); err != nil {
		// If the error is other than not found, abort immediately.
		if err != dockerclient.ErrNotFound || !pullImage {
			return nil, err
		}
		// Otherwise, try to pull the image...
		if err = n.Pull(config.Image); err != nil {
			return nil, err
		}
		// ...And try again.
		if id, err = client.CreateContainer(&newConfig, name); err != nil {
			return nil, err
		}
	}

	// Register the container immediately while waiting for a state refresh.
	// Force a state refresh to pick up the newly created container.
	n.RefreshContainer(id, true)

	n.RLock()
	defer n.RUnlock()

	return n.containers[id], nil
}

// Destroy and remove a container from the node.
func (n *Node) Destroy(container *Container, force bool) error {
	if err := n.client.RemoveContainer(container.Id, force); err != nil {
		return err
	}

	// Remove the container from the state. Eventually, the state refresh loop
	// will rewrite this.
	n.Lock()
	defer n.Unlock()
	delete(n.containers, container.Id)

	return nil
}

func (n *Node) Pull(image string) error {
	if err := n.client.PullImage(image, nil); err != nil {
		return err
	}
	return nil
}

// Register an event handler.
func (n *Node) Events(h EventHandler) error {
	if n.eventHandler != nil {
		return errors.New("event handler already set")
	}
	n.eventHandler = h
	return nil
}

func (n *Node) Containers() []*Container {
	containers := []*Container{}
	n.RLock()
	for _, container := range n.containers {
		containers = append(containers, container)
	}
	n.RUnlock()
	return containers
}

func (n *Node) Images() []*dockerclient.Image {
	images := []*dockerclient.Image{}
	n.RLock()
	for _, image := range n.images {
		images = append(images, image)
	}
	n.RUnlock()
	return images
}

// Image returns the image with IdOrName in the node
func (n *Node) Image(IdOrName string) *dockerclient.Image {
	size := len(IdOrName)
	for _, image := range n.Images() {
		if image.Id == IdOrName || (size > 2 && strings.HasPrefix(image.Id, IdOrName)) {
			return image
		}
		for _, t := range image.RepoTags {
			if t == IdOrName || (size > 2 && strings.HasPrefix(t, IdOrName)) {
				return image
			}
		}
	}
	return nil
}

func (n *Node) String() string {
	return fmt.Sprintf("node %s addr %s", n.ID, n.Addr)
}

func (n *Node) handler(ev *dockerclient.Event, _ chan error, args ...interface{}) {
	// Something changed - refresh our internal state.
	switch ev.Status {
	case "pull", "untag", "delete":
		// These events refer to images so there's no need to update
		// containers.
		n.refreshImages()
	case "start", "die":
		// If the container is started or stopped, we have to do an inspect in
		// order to get the new NetworkSettings.
		n.RefreshContainer(ev.Id, true)
	default:
		// Otherwise, do a "soft" refresh of the container.
		n.RefreshContainer(ev.Id, false)
	}

	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}

	event := &Event{
		Node:  n,
		Event: *ev,
	}

	n.eventHandler.Handle(event)
}

// Inject a container into the internal state.
func (n *Node) AddContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; ok {
		return errors.New("container already exists")
	}
	n.containers[container.Id] = container
	return nil
}

// Inject an image into the internal state.
func (n *Node) AddImage(image *dockerclient.Image) {
	n.Lock()
	defer n.Unlock()

	n.images = append(n.images, image)
}

// Remove a container from the internal test.
func (n *Node) RemoveContainer(container *Container) error {
	n.Lock()
	defer n.Unlock()

	if _, ok := n.containers[container.Id]; !ok {
		return errors.New("container not found")
	}
	delete(n.containers, container.Id)
	return nil
}

// Wipes the internal container state.
func (n *Node) CleanupContainers() {
	n.Lock()
	n.containers = make(map[string]*Container)
	n.Unlock()
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Pulcy.
//
// 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 main

import (
	"fmt"

	"github.com/ryanuber/columnize"
	"github.com/spf13/cobra"

	"github.com/pulcy/quark/providers"
)

var (
	cmdClusterInfo = &cobra.Command{
		Use: "info",
		Run: showClusterInfo,
	}

	clusterInfoFlags providers.ClusterInfo
)

func init() {
	cmdClusterInfo.Flags().StringVar(&clusterInfoFlags.Domain, "domain", defaultDomain(), "Cluster domain")
	cmdClusterInfo.Flags().StringVar(&clusterInfoFlags.Name, "name", "", "Cluster name")
	cmdCluster.AddCommand(cmdClusterInfo)
}

func showClusterInfo(cmd *cobra.Command, args []string) {
	clusterInfoFromArgs(&clusterInfoFlags, args)

	provider := newProvider()
	clusterInfoFlags = provider.ClusterDefaults(clusterInfoFlags)

	if clusterInfoFlags.Name == "" {
		Exitf("Please specify a name\n")
	}
	instances, err := provider.GetInstances(instancesFlags)
	if err != nil {
		Exitf("Failed to list instances: %v\n", err)
	}
	clusterMembers, err := instances.AsClusterMemberList(log, nil)
	if err != nil {
		Exitf("Failed to fetch instance member data: %v\n", err)
	}

	lines := []string{
		fmt.Sprintf("ID | %s", clusterMembers[0].ClusterID),
		fmt.Sprintf("#Instances | %d", len(instances)),
	}
	result := columnize.SimpleFormat(lines)
	fmt.Println(result)
}
<commit_msg>Fixed `cluster info` command<commit_after>// Copyright (c) 2016 Pulcy.
//
// 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 main

import (
	"fmt"

	"github.com/ryanuber/columnize"
	"github.com/spf13/cobra"

	"github.com/pulcy/quark/providers"
)

var (
	cmdClusterInfo = &cobra.Command{
		Use: "info",
		Run: showClusterInfo,
	}

	clusterInfoFlags providers.ClusterInfo
)

func init() {
	cmdClusterInfo.Flags().StringVar(&clusterInfoFlags.Domain, "domain", defaultDomain(), "Cluster domain")
	cmdClusterInfo.Flags().StringVar(&clusterInfoFlags.Name, "name", "", "Cluster name")
	cmdCluster.AddCommand(cmdClusterInfo)
}

func showClusterInfo(cmd *cobra.Command, args []string) {
	clusterInfoFromArgs(&clusterInfoFlags, args)

	provider := newProvider()
	clusterInfoFlags = provider.ClusterDefaults(clusterInfoFlags)

	if clusterInfoFlags.Name == "" {
		Exitf("Please specify a name\n")
	}
	instances, err := provider.GetInstances(clusterInfoFlags)
	if err != nil {
		Exitf("Failed to list instances: %v\n", err)
	}
	clusterMembers, err := instances.AsClusterMemberList(log, nil)
	if err != nil {
		Exitf("Failed to fetch instance member data: %v\n", err)
	}

	lines := []string{
		fmt.Sprintf("ID | %s", clusterMembers[0].ClusterID),
		fmt.Sprintf("#Instances | %d", len(instances)),
	}
	result := columnize.SimpleFormat(lines)
	fmt.Println(result)
}
<|endoftext|>


<commit_before>package main

import (
	"bla"
	"flag"
	"log"
	"net/http"
	"sync"
	"time"
)

const (
	DefaultConfig = "config.json"
)

var (
	certfile   = flag.String("cert", "", "cert file path")
	keyfile    = flag.String("key", "", "cert file path")
	addr       = flag.String("addr", ":8080", "listen port")
	configPath = flag.String("config", DefaultConfig, "default config path")
	logPool    = sync.Pool{New: func() interface{} { return &LogWriter{nil, 200} }}
)

func main() {

	flag.Parse()
	h := bla.NewHandler(*configPath)
	log.Print("addr: ", *addr)
	log.Print("cfg: ", *configPath)

	lh := logTimeAndStatus(h)

	if *certfile != "" && *keyfile != "" {
		log.Printf("TLS:%s, %s", *certfile, *keyfile)
		http.ListenAndServeTLS(*addr, *certfile, *keyfile, lh)
		return
	}
	http.ListenAndServe(*addr, lh)

}

func logTimeAndStatus(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		writer := logPool.Get().(*LogWriter)
		writer.ResponseWriter = w
		writer.statusCode = 200

		handler.ServeHTTP(writer, r)
		log.Printf("%s %s %s %s %d", r.RemoteAddr,
			r.Method, r.URL.Path, time.Now().Sub(start), writer.statusCode)
		logPool.Put(writer)
	})
}

type LogWriter struct {
	http.ResponseWriter
	statusCode int
}

func (l *LogWriter) WriteHeader(i int) {
	l.statusCode = i
	l.ResponseWriter.WriteHeader(i)
}
<commit_msg>update cmd<commit_after>package main

import (
	"flag"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/mengzhuo/bla"
)

const (
	DefaultConfig = "config.json"
)

var (
	certfile   = flag.String("cert", "", "cert file path")
	keyfile    = flag.String("key", "", "cert file path")
	addr       = flag.String("addr", ":8080", "listen port")
	configPath = flag.String("config", DefaultConfig, "default config path")
	logPool    = sync.Pool{New: func() interface{} { return &LogWriter{nil, 200} }}
)

func main() {

	flag.Parse()
	h := bla.NewHandler(*configPath)
	log.Print("addr: ", *addr)
	log.Print("cfg: ", *configPath)

	lh := logTimeAndStatus(h)

	if *certfile != "" && *keyfile != "" {
		log.Printf("TLS:%s, %s", *certfile, *keyfile)
		http.ListenAndServeTLS(*addr, *certfile, *keyfile, lh)
		return
	}
	http.ListenAndServe(*addr, lh)

}

func logTimeAndStatus(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		writer := logPool.Get().(*LogWriter)
		writer.ResponseWriter = w
		writer.statusCode = 200

		handler.ServeHTTP(writer, r)
		log.Printf("%s %s %s %s %d", r.RemoteAddr,
			r.Method, r.URL.Path, time.Now().Sub(start), writer.statusCode)
		logPool.Put(writer)
	})
}

type LogWriter struct {
	http.ResponseWriter
	statusCode int
}

func (l *LogWriter) WriteHeader(i int) {
	l.statusCode = i
	l.ResponseWriter.WriteHeader(i)
}
<|endoftext|>


<commit_before>//  Copyright (c) 2015 Couchbase, Inc.
//  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 main

import (
	"bytes"
	"expvar"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"path"
	"runtime"
	"runtime/pprof"
	"sort"
	"strings"
	"time"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/cbgt"
	"github.com/couchbaselabs/cbgt/cmd"
	"github.com/couchbaselabs/cbgt/rebalance"
)

var VERSION = "v0.0.0"

var expvars = expvar.NewMap("stats")

func main() {
	flag.Parse()

	if flags.Help {
		flag.Usage()
		os.Exit(2)
	}

	if flags.Version {
		fmt.Printf("%s main: %s, data: %s\n",
			path.Base(os.Args[0]), VERSION, cbgt.VERSION)
		os.Exit(0)
	}

	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	log.Printf("main: %s started (%s/%s)",
		os.Args[0], VERSION, cbgt.VERSION)

	rand.Seed(time.Now().UTC().UnixNano())

	go dumpOnSignalForPlatform()

	MainWelcome(flagAliases)

	nodesToRemove := []string(nil)
	if len(flags.RemoveNodes) > 0 {
		nodesToRemove = strings.Split(flags.RemoveNodes, ",")
	}

	bindHttp := "NO-BIND-HTTP"
	register := "unchanged"
	dataDir := "NO-DATA-DIR"

	// If cfg is down, we error, leaving it to some user-supplied
	// outside watchdog to backoff and restart/retry.
	cfg, err := cmd.MainCfg("mcp", flags.CfgConnect,
		bindHttp, register, dataDir)
	if err != nil {
		if err == cmd.ErrorBindHttp {
			log.Fatalf("%v", err)
			return
		}
		log.Fatalf("main: could not start cfg, cfgConnect: %s, err: %v\n"+
			"  Please check that your -cfg/-cfgConnect parameter (%q)\n"+
			"  is correct and/or that your configuration provider\n"+
			"  is available.",
			flags.CfgConnect, err, flags.CfgConnect)
		return
	}

	r, err := rebalance.StartRebalance(cbgt.VERSION, cfg, flags.Server,
		nodesToRemove,
		rebalance.RebalanceOptions{
			DryRun:  flags.DryRun,
			Verbose: flags.Verbose,
		})
	if err != nil {
		log.Fatalf("main: StartRebalance, err: %v", err)
		return
	}

	reportProgress(r)

	r.Stop()

	log.Printf("main: done")
}

// ------------------------------------------------------------

type ProgressEntry struct {
	pindex, sourcePartition, node string // Immutable.

	stateOp     rebalance.StateOp
	initUUIDSeq cbgt.UUIDSeq
	currUUIDSeq cbgt.UUIDSeq
	wantUUIDSeq cbgt.UUIDSeq
}

func reportProgress(r *rebalance.Rebalancer) {
	maxNodeLen := 0
	maxPIndexLen := 0

	seenNodes := map[string]bool{}
	seenNodesSorted := []string(nil)

	// Map of pindex -> (source) partition -> node -> *ProgressEntry
	progressEntries := map[string]map[string]map[string]*ProgressEntry{}

	inflightPIndexes := map[string]bool{}
	inflightPIndexesSorted := []string(nil)

	updateProgressEntry := func(pindex, sourcePartition, node string,
		cb func(*ProgressEntry)) {
		if !seenNodes[node] {
			seenNodes[node] = true
			seenNodesSorted = append(seenNodesSorted, node)
			sort.Strings(seenNodesSorted)

			if maxNodeLen < len(node) {
				maxNodeLen = len(node)
			}
		}

		if maxPIndexLen < len(pindex) {
			maxPIndexLen = len(pindex)
		}

		sourcePartitions, exists := progressEntries[pindex]
		if !exists || sourcePartitions == nil {
			sourcePartitions = map[string]map[string]*ProgressEntry{}
			progressEntries[pindex] = sourcePartitions
		}

		nodes, exists := sourcePartitions[sourcePartition]
		if !exists || nodes == nil {
			nodes = map[string]*ProgressEntry{}
			sourcePartitions[sourcePartition] = nodes
		}

		progressEntry, exists := nodes[node]
		if !exists || progressEntry == nil {
			progressEntry = &ProgressEntry{
				pindex:          pindex,
				sourcePartition: sourcePartition,
				node:            node,
			}
			nodes[node] = progressEntry
		}

		cb(progressEntry)

		// TODO: Check UUID matches, too.

		if progressEntry.wantUUIDSeq.Seq >
			progressEntry.currUUIDSeq.Seq {
			if !inflightPIndexes[pindex] {
				inflightPIndexes[pindex] = true
				inflightPIndexesSorted =
					append(inflightPIndexesSorted, pindex)

				sort.Strings(inflightPIndexesSorted)
			}
		}
	}

	for progress := range r.ProgressCh() {
		if progress.Index == "" {
			r.Log("main: progress: %+v", progress)
			continue
		}

		if progress.Error != nil {
			r.Log("main: error, progress: %+v", progress)
			continue
		}

		r.Visit(func(
			currStates rebalance.CurrStates,
			currSeqs rebalance.CurrSeqs,
			wantSeqs rebalance.WantSeqs) {
			for _, pindexes := range currStates {
				for pindex, nodes := range pindexes {
					for node, stateOp := range nodes {
						updateProgressEntry(pindex, "", node,
							func(pe *ProgressEntry) {
								pe.stateOp = stateOp
							})
					}
				}
			}

			for pindex, sourcePartitions := range currSeqs {
				for sourcePartition, nodes := range sourcePartitions {
					for node, currUUIDSeq := range nodes {
						updateProgressEntry(pindex,
							sourcePartition, node,
							func(pe *ProgressEntry) {
								pe.currUUIDSeq = currUUIDSeq

								if pe.initUUIDSeq.UUID == "" {
									pe.initUUIDSeq = currUUIDSeq
								}
							})
					}
				}
			}

			for pindex, sourcePartitions := range wantSeqs {
				for sourcePartition, nodes := range sourcePartitions {
					for node, wantUUIDSeq := range nodes {
						updateProgressEntry(pindex,
							sourcePartition, node,
							func(pe *ProgressEntry) {
								pe.wantUUIDSeq = wantUUIDSeq
							})
					}
				}
			}

			// ----------------------------------------

			var b bytes.Buffer

			for i := 0; i < maxPIndexLen; i++ {
				b.WriteByte(' ')
			}
			b.WriteByte(' ')

			for i, seenNode := range seenNodesSorted {
				if i > 0 {
					b.WriteByte(' ')
				}
				b.Write([]byte(seenNode))
			}
			b.WriteByte('\n')

			for _, inflightPIndex := range inflightPIndexesSorted {
				b.Write([]byte("                    "))
				b.Write([]byte(inflightPIndex))

				for _, seenNode := range seenNodesSorted {
					b.WriteByte(' ')

					sourcePartitions, exists :=
						progressEntries[inflightPIndex]
					if !exists || sourcePartitions == nil {
						emitNodeEntry(&b, nil, nil, maxNodeLen)
						continue
					}

					nodes, exists := sourcePartitions[""]
					if !exists || nodes == nil {
						emitNodeEntry(&b, nil, nil, maxNodeLen)
						continue
					}

					pe, exists := nodes[seenNode]
					if !exists || pe == nil {
						emitNodeEntry(&b, nil, nil, maxNodeLen)
						continue
					}

					emitNodeEntry(&b, pe, sourcePartitions, maxNodeLen)
				}

				b.WriteByte('\n')
			}

			r.Log("%s", b.String())
		})
	}
}

var opMap = map[string]string{
	"":        ".",
	"add":     "+",
	"del":     "-",
	"promote": "P",
	"demote":  "D",
}

func emitNodeEntry(b *bytes.Buffer,
	pe *ProgressEntry,
	sourcePartitions map[string]map[string]*ProgressEntry,
	maxNodeLen int) {
	written := 0

	totPct := 0.0 // To compute average pct.
	numPct := 0

	if pe != nil &&
		sourcePartitions != nil {
		written, _ = b.Write([]byte(opMap[pe.stateOp.Op]))

		for sourcePartition, nodes := range sourcePartitions {
			if sourcePartition == "" {
				continue
			}

			pex := nodes[pe.node]
			if pex == nil || pex.wantUUIDSeq.UUID == "" {
				continue
			}

			if pex.wantUUIDSeq.Seq <= pex.currUUIDSeq.Seq {
				totPct = totPct + 1.0
				numPct = numPct + 1
				continue
			}

			n := pex.currUUIDSeq.Seq - pex.initUUIDSeq.Seq
			d := pex.wantUUIDSeq.Seq - pex.initUUIDSeq.Seq
			if d > 0 {
				pct := float64(n) / float64(d)
				totPct = totPct + pct
				numPct = numPct + 1
			}
		}
	} else {
		b.WriteByte('.')
		written = 1
	}

	if numPct > 0 {
		avgPct := totPct / float64(numPct)

		n, _ := fmt.Fprintf(b, " %.1f%%", avgPct*100.0)
		written = written + n
	}

	for i := written; i < maxNodeLen; i++ {
		b.WriteByte(' ')
	}
}

// ------------------------------------------------------------

func MainWelcome(flagAliases map[string][]string) {
	flag.VisitAll(func(f *flag.Flag) {
		if flagAliases[f.Name] != nil {
			log.Printf("  -%s=%q\n", f.Name, f.Value)
		}
	})
	log.Printf("  GOMAXPROCS=%d", runtime.GOMAXPROCS(-1))
}

func dumpOnSignal(signals ...os.Signal) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, signals...)
	for _ = range c {
		log.Printf("dump: goroutine...")
		pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
		log.Printf("dump: heap...")
		pprof.Lookup("heap").WriteTo(os.Stderr, 1)
	}
}
<commit_msg>looser check for inflightPIndexes<commit_after>//  Copyright (c) 2015 Couchbase, Inc.
//  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 main

import (
	"bytes"
	"expvar"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"path"
	"runtime"
	"runtime/pprof"
	"sort"
	"strings"
	"time"

	log "github.com/couchbase/clog"

	"github.com/couchbaselabs/cbgt"
	"github.com/couchbaselabs/cbgt/cmd"
	"github.com/couchbaselabs/cbgt/rebalance"
)

var VERSION = "v0.0.0"

var expvars = expvar.NewMap("stats")

func main() {
	flag.Parse()

	if flags.Help {
		flag.Usage()
		os.Exit(2)
	}

	if flags.Version {
		fmt.Printf("%s main: %s, data: %s\n",
			path.Base(os.Args[0]), VERSION, cbgt.VERSION)
		os.Exit(0)
	}

	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	log.Printf("main: %s started (%s/%s)",
		os.Args[0], VERSION, cbgt.VERSION)

	rand.Seed(time.Now().UTC().UnixNano())

	go dumpOnSignalForPlatform()

	MainWelcome(flagAliases)

	nodesToRemove := []string(nil)
	if len(flags.RemoveNodes) > 0 {
		nodesToRemove = strings.Split(flags.RemoveNodes, ",")
	}

	bindHttp := "NO-BIND-HTTP"
	register := "unchanged"
	dataDir := "NO-DATA-DIR"

	// If cfg is down, we error, leaving it to some user-supplied
	// outside watchdog to backoff and restart/retry.
	cfg, err := cmd.MainCfg("mcp", flags.CfgConnect,
		bindHttp, register, dataDir)
	if err != nil {
		if err == cmd.ErrorBindHttp {
			log.Fatalf("%v", err)
			return
		}
		log.Fatalf("main: could not start cfg, cfgConnect: %s, err: %v\n"+
			"  Please check that your -cfg/-cfgConnect parameter (%q)\n"+
			"  is correct and/or that your configuration provider\n"+
			"  is available.",
			flags.CfgConnect, err, flags.CfgConnect)
		return
	}

	r, err := rebalance.StartRebalance(cbgt.VERSION, cfg, flags.Server,
		nodesToRemove,
		rebalance.RebalanceOptions{
			DryRun:  flags.DryRun,
			Verbose: flags.Verbose,
		})
	if err != nil {
		log.Fatalf("main: StartRebalance, err: %v", err)
		return
	}

	reportProgress(r)

	r.Stop()

	log.Printf("main: done")
}

// ------------------------------------------------------------

type ProgressEntry struct {
	pindex, sourcePartition, node string // Immutable.

	stateOp     rebalance.StateOp
	initUUIDSeq cbgt.UUIDSeq
	currUUIDSeq cbgt.UUIDSeq
	wantUUIDSeq cbgt.UUIDSeq
}

func reportProgress(r *rebalance.Rebalancer) {
	maxNodeLen := 0
	maxPIndexLen := 0

	seenNodes := map[string]bool{}
	seenNodesSorted := []string(nil)

	// Map of pindex -> (source) partition -> node -> *ProgressEntry
	progressEntries := map[string]map[string]map[string]*ProgressEntry{}

	inflightPIndexes := map[string]bool{}
	inflightPIndexesSorted := []string(nil)

	updateProgressEntry := func(pindex, sourcePartition, node string,
		cb func(*ProgressEntry)) {
		if !seenNodes[node] {
			seenNodes[node] = true
			seenNodesSorted = append(seenNodesSorted, node)
			sort.Strings(seenNodesSorted)

			if maxNodeLen < len(node) {
				maxNodeLen = len(node)
			}
		}

		if maxPIndexLen < len(pindex) {
			maxPIndexLen = len(pindex)
		}

		sourcePartitions, exists := progressEntries[pindex]
		if !exists || sourcePartitions == nil {
			sourcePartitions = map[string]map[string]*ProgressEntry{}
			progressEntries[pindex] = sourcePartitions
		}

		nodes, exists := sourcePartitions[sourcePartition]
		if !exists || nodes == nil {
			nodes = map[string]*ProgressEntry{}
			sourcePartitions[sourcePartition] = nodes
		}

		progressEntry, exists := nodes[node]
		if !exists || progressEntry == nil {
			progressEntry = &ProgressEntry{
				pindex:          pindex,
				sourcePartition: sourcePartition,
				node:            node,
			}
			nodes[node] = progressEntry
		}

		cb(progressEntry)

		// TODO: Check UUID matches, too.

		if !inflightPIndexes[pindex] {
			inflightPIndexes[pindex] = true
			inflightPIndexesSorted =
				append(inflightPIndexesSorted, pindex)

			sort.Strings(inflightPIndexesSorted)
		}
	}

	for progress := range r.ProgressCh() {
		if progress.Index == "" {
			r.Log("main: progress: %+v", progress)
			continue
		}

		if progress.Error != nil {
			r.Log("main: error, progress: %+v", progress)
			continue
		}

		r.Visit(func(
			currStates rebalance.CurrStates,
			currSeqs rebalance.CurrSeqs,
			wantSeqs rebalance.WantSeqs) {
			for _, pindexes := range currStates {
				for pindex, nodes := range pindexes {
					for node, stateOp := range nodes {
						updateProgressEntry(pindex, "", node,
							func(pe *ProgressEntry) {
								pe.stateOp = stateOp
							})
					}
				}
			}

			for pindex, sourcePartitions := range currSeqs {
				for sourcePartition, nodes := range sourcePartitions {
					for node, currUUIDSeq := range nodes {
						updateProgressEntry(pindex,
							sourcePartition, node,
							func(pe *ProgressEntry) {
								pe.currUUIDSeq = currUUIDSeq

								if pe.initUUIDSeq.UUID == "" {
									pe.initUUIDSeq = currUUIDSeq
								}
							})
					}
				}
			}

			for pindex, sourcePartitions := range wantSeqs {
				for sourcePartition, nodes := range sourcePartitions {
					for node, wantUUIDSeq := range nodes {
						updateProgressEntry(pindex,
							sourcePartition, node,
							func(pe *ProgressEntry) {
								pe.wantUUIDSeq = wantUUIDSeq
							})
					}
				}
			}

			// ----------------------------------------

			var b bytes.Buffer

			for i := 0; i < maxPIndexLen; i++ {
				b.WriteByte(' ')
			}
			b.WriteByte(' ')

			for i, seenNode := range seenNodesSorted {
				if i > 0 {
					b.WriteByte(' ')
				}
				b.Write([]byte(seenNode))
			}
			b.WriteByte('\n')

			for _, inflightPIndex := range inflightPIndexesSorted {
				b.Write([]byte("                    "))
				b.Write([]byte(inflightPIndex))

				for _, seenNode := range seenNodesSorted {
					b.WriteByte(' ')

					sourcePartitions, exists :=
						progressEntries[inflightPIndex]
					if !exists || sourcePartitions == nil {
						emitNodeEntry(&b, nil, nil, maxNodeLen)
						continue
					}

					nodes, exists := sourcePartitions[""]
					if !exists || nodes == nil {
						emitNodeEntry(&b, nil, nil, maxNodeLen)
						continue
					}

					pe, exists := nodes[seenNode]
					if !exists || pe == nil {
						emitNodeEntry(&b, nil, nil, maxNodeLen)
						continue
					}

					emitNodeEntry(&b, pe, sourcePartitions, maxNodeLen)
				}

				b.WriteByte('\n')
			}

			r.Log("%s", b.String())
		})
	}
}

var opMap = map[string]string{
	"":        ".",
	"add":     "+",
	"del":     "-",
	"promote": "P",
	"demote":  "D",
}

func emitNodeEntry(b *bytes.Buffer,
	pe *ProgressEntry,
	sourcePartitions map[string]map[string]*ProgressEntry,
	maxNodeLen int) {
	written := 0

	totPct := 0.0 // To compute average pct.
	numPct := 0

	if pe != nil &&
		sourcePartitions != nil {
		written, _ = b.Write([]byte(opMap[pe.stateOp.Op]))

		for sourcePartition, nodes := range sourcePartitions {
			if sourcePartition == "" {
				continue
			}

			pex := nodes[pe.node]
			if pex == nil || pex.wantUUIDSeq.UUID == "" {
				continue
			}

			if pex.wantUUIDSeq.Seq <= pex.currUUIDSeq.Seq {
				totPct = totPct + 1.0
				numPct = numPct + 1
				continue
			}

			n := pex.currUUIDSeq.Seq - pex.initUUIDSeq.Seq
			d := pex.wantUUIDSeq.Seq - pex.initUUIDSeq.Seq
			if d > 0 {
				pct := float64(n) / float64(d)
				totPct = totPct + pct
				numPct = numPct + 1
			}
		}
	} else {
		b.WriteByte('.')
		written = 1
	}

	if numPct > 0 {
		avgPct := totPct / float64(numPct)

		n, _ := fmt.Fprintf(b, " %.1f%%", avgPct*100.0)
		written = written + n
	}

	for i := written; i < maxNodeLen; i++ {
		b.WriteByte(' ')
	}
}

// ------------------------------------------------------------

func MainWelcome(flagAliases map[string][]string) {
	flag.VisitAll(func(f *flag.Flag) {
		if flagAliases[f.Name] != nil {
			log.Printf("  -%s=%q\n", f.Name, f.Value)
		}
	})
	log.Printf("  GOMAXPROCS=%d", runtime.GOMAXPROCS(-1))
}

func dumpOnSignal(signals ...os.Signal) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, signals...)
	for _ = range c {
		log.Printf("dump: goroutine...")
		pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
		log.Printf("dump: heap...")
		pprof.Lookup("heap").WriteTo(os.Stderr, 1)
	}
}
<|endoftext|>