<commit_before>// run_chromium_perf_on_workers is an application that runs the specified telemetry
// benchmark on all CT workers and uploads the results to Google Storage. The
// requester is emailed when the task is done.
package main

import (
	"database/sql"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/skia-dev/glog"
	"go.skia.org/infra/ct/go/ctfe/chromium_perf"
	"go.skia.org/infra/ct/go/frontend"
	"go.skia.org/infra/ct/go/master_scripts/master_common"
	"go.skia.org/infra/ct/go/util"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/email"
	skutil "go.skia.org/infra/go/util"
)

const (
	MAX_PAGES_PER_SWARMING_BOT = 100
)

var (
	emails                    = flag.String("emails", "", "The comma separated email addresses to notify when the task is picked up and completes.")
	description               = flag.String("description", "", "The description of the run as entered by the requester.")
	gaeTaskID                 = flag.Int64("gae_task_id", -1, "The key of the App Engine task. This task will be updated when the task is completed.")
	pagesetType               = flag.String("pageset_type", "", "The type of pagesets to use. Eg: 10k, Mobile10k, All.")
	benchmarkName             = flag.String("benchmark_name", "", "The telemetry benchmark to run on the workers.")
	benchmarkExtraArgs        = flag.String("benchmark_extra_args", "", "The extra arguments that are passed to the specified benchmark.")
	browserExtraArgsNoPatch   = flag.String("browser_extra_args_nopatch", "", "The extra arguments that are passed to the browser while running the benchmark for the nopatch case.")
	browserExtraArgsWithPatch = flag.String("browser_extra_args_withpatch", "", "The extra arguments that are passed to the browser while running the benchmark for the withpatch case.")
	repeatBenchmark           = flag.Int("repeat_benchmark", 1, "The number of times the benchmark should be repeated. For skpicture_printer benchmark this value is always 1.")
	runInParallel             = flag.Bool("run_in_parallel", false, "Run the benchmark by bringing up multiple chrome instances in parallel.")
	targetPlatform            = flag.String("target_platform", util.PLATFORM_ANDROID, "The platform the benchmark will run on (Android / Linux).")
	runID                     = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
	varianceThreshold         = flag.Float64("variance_threshold", 0.0, "The variance threshold to use when comparing the resultant CSV files.")
	discardOutliers           = flag.Float64("discard_outliers", 0.0, "The percentage of outliers to discard when comparing the result CSV files.")

	taskCompletedSuccessfully = false

	htmlOutputLink      = util.MASTER_LOGSERVER_LINK
	skiaPatchLink       = util.MASTER_LOGSERVER_LINK
	chromiumPatchLink   = util.MASTER_LOGSERVER_LINK
	catapultPatchLink   = util.MASTER_LOGSERVER_LINK
	benchmarkPatchLink  = util.MASTER_LOGSERVER_LINK
	noPatchOutputLink   = util.MASTER_LOGSERVER_LINK
	withPatchOutputLink = util.MASTER_LOGSERVER_LINK
)

func sendEmail(recipients []string) {
	// Send completion email.
	emailSubject := fmt.Sprintf("Cluster telemetry chromium perf task has completed (%s)", *runID)
	failureHtml := ""
	viewActionMarkup := ""
	var err error

	if taskCompletedSuccessfully {
		if viewActionMarkup, err = email.GetViewActionMarkup(htmlOutputLink, "View Results", "Direct link to the HTML results"); err != nil {
			glog.Errorf("Failed to get view action markup: %s", err)
			return
		}
	} else {
		emailSubject += " with failures"
		failureHtml = util.GetFailureEmailHtml(*runID)
		if viewActionMarkup, err = email.GetViewActionMarkup(util.GetMasterLogLink(*runID), "View Failure", "Direct link to the master log"); err != nil {
			glog.Errorf("Failed to get view action markup: %s", err)
			return
		}
	}
	bodyTemplate := `
	The chromium perf %s benchmark task on %s pageset has completed.

	Run description: %s

	%s
	The HTML output with differences between the base run and the patch run is here.

	The patch(es) you specified are here:
	chromium/skia/catapult

	You can schedule more runs here.

	Thanks!
	`
	emailBody := fmt.Sprintf(bodyTemplate, *benchmarkName, *pagesetType, *description, failureHtml, htmlOutputLink, chromiumPatchLink, skiaPatchLink, catapultPatchLink, frontend.ChromiumPerfTasksWebapp)
	if err := util.SendEmailWithMarkup(recipients, emailSubject, emailBody, viewActionMarkup); err != nil {
		glog.Errorf("Error while sending email: %s", err)
		return
	}
}

func updateWebappTask() {
	vars := chromium_perf.UpdateVars{}
	vars.Id = *gaeTaskID
	vars.SetCompleted(taskCompletedSuccessfully)
	vars.Results = sql.NullString{String: htmlOutputLink, Valid: true}
	vars.NoPatchRawOutput = sql.NullString{String: noPatchOutputLink, Valid: true}
	vars.WithPatchRawOutput = sql.NullString{String: withPatchOutputLink, Valid: true}
	skutil.LogErr(frontend.UpdateWebappTaskV2(&vars))
}

func main() {
	defer common.LogPanic()
	master_common.Init()

	// Send start email.
	emailsArr := util.ParseEmails(*emails)
	emailsArr = append(emailsArr, util.CtAdmins...)
	if len(emailsArr) == 0 {
		glog.Error("At least one email address must be specified")
		return
	}
	skutil.LogErr(frontend.UpdateWebappTaskSetStarted(&chromium_perf.UpdateVars{}, *gaeTaskID))
	skutil.LogErr(util.SendTaskStartEmail(emailsArr, "Chromium perf", *runID, *description))
	// Ensure webapp is updated and email is sent even if task fails.
	defer updateWebappTask()
	defer sendEmail(emailsArr)
	// Cleanup dirs after run completes.
	defer skutil.RemoveAll(filepath.Join(util.StorageDir, util.ChromiumPerfRunsDir, *runID))
	// Finish with glog flush and how long the task took.
	defer util.TimeTrack(time.Now(), "Running chromium perf task on workers")
	defer glog.Flush()

	if *pagesetType == "" {
		glog.Error("Must specify --pageset_type")
		return
	}
	if *benchmarkName == "" {
		glog.Error("Must specify --benchmark_name")
		return
	}
	if *runID == "" {
		glog.Error("Must specify --run_id")
		return
	}
	if *description == "" {
		glog.Error("Must specify --description")
		return
	}

	// Instantiate GsUtil object.
	gs, err := util.NewGsUtil(nil)
	if err != nil {
		glog.Errorf("Could not instantiate gsutil object: %s", err)
		return
	}
	remoteOutputDir := filepath.Join(util.ChromiumPerfRunsDir, *runID)

	// Copy the patches to Google Storage.
	skiaPatchName := *runID + ".skia.patch"
	chromiumPatchName := *runID + ".chromium.patch"
	catapultPatchName := *runID + ".catapult.patch"
	benchmarkPatchName := *runID + ".benchmark.patch"
	for _, patchName := range []string{skiaPatchName, chromiumPatchName, catapultPatchName, benchmarkPatchName} {
		if err := gs.UploadFile(patchName, os.TempDir(), remoteOutputDir); err != nil {
			glog.Errorf("Could not upload %s to %s: %s", patchName, remoteOutputDir, err)
			return
		}
	}
	skiaPatchLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, remoteOutputDir, skiaPatchName)
	chromiumPatchLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, remoteOutputDir, chromiumPatchName)
	catapultPatchLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, remoteOutputDir, catapultPatchName)
	benchmarkPatchLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, remoteOutputDir, benchmarkPatchName)

	// Create the two required chromium builds (with patch and without the patch).
	chromiumBuilds, err := util.TriggerBuildRepoSwarmingTask(
		"build_chromium", *runID, "chromium", *targetPlatform, []string{},
		[]string{filepath.Join(remoteOutputDir, chromiumPatchName), filepath.Join(remoteOutputDir, skiaPatchName)},
		/*singlebuild*/ false, 3*time.Hour, 1*time.Hour)
	if err != nil {
		glog.Errorf("Error encountered when swarming build repo task: %s", err)
		return
	}
	if len(chromiumBuilds) != 2 {
		glog.Errorf("Expected 2 builds but instead got %d: %v.", len(chromiumBuilds), chromiumBuilds)
		return
	}
	chromiumBuildNoPatch := chromiumBuilds[0]
	chromiumBuildWithPatch := chromiumBuilds[1]

	// Parse out the Chromium and Skia hashes.
	chromiumHash, skiaHash := getHashesFromBuild(chromiumBuildNoPatch)

	// Archive, trigger and collect swarming tasks.
	isolateExtraArgs := map[string]string{
		"CHROMIUM_BUILD_NOPATCH":       chromiumBuildNoPatch,
		"CHROMIUM_BUILD_WITHPATCH":     chromiumBuildWithPatch,
		"RUN_ID":                       *runID,
		"BENCHMARK":                    *benchmarkName,
		"BENCHMARK_ARGS":               *benchmarkExtraArgs,
		"BROWSER_EXTRA_ARGS_NOPATCH":   *browserExtraArgsNoPatch,
		"BROWSER_EXTRA_ARGS_WITHPATCH": *browserExtraArgsWithPatch,
		"REPEAT_BENCHMARK":             strconv.Itoa(*repeatBenchmark),
		"RUN_IN_PARALLEL":              strconv.FormatBool(*runInParallel),
		"TARGET_PLATFORM":              *targetPlatform,
	}
	if err := util.TriggerSwarmingTask(*pagesetType, "chromium_perf", util.CHROMIUM_PERF_ISOLATE, *runID, 12*time.Hour, 1*time.Hour, util.USER_TASKS_PRIORITY, MAX_PAGES_PER_SWARMING_BOT, isolateExtraArgs, util.GOLO_WORKER_DIMENSIONS); err != nil {
		glog.Errorf("Error encountered when swarming tasks: %s", err)
		return
	}

	// If "--output-format=csv-pivot-table" was specified then merge all CSV files and upload.
	runIDNoPatch := fmt.Sprintf("%s-nopatch", *runID)
	runIDWithPatch := fmt.Sprintf("%s-withpatch", *runID)
	noOutputSlaves := []string{}
	pathToPyFiles := util.GetPathToPyFiles(false)
	for _, run := range []string{runIDNoPatch, runIDWithPatch} {
		if strings.Contains(*benchmarkExtraArgs, "--output-format=csv-pivot-table") {
			if noOutputSlaves, err = util.MergeUploadCSVFiles(run, pathToPyFiles, gs, util.PagesetTypeToInfo[*pagesetType].NumPages, MAX_PAGES_PER_SWARMING_BOT, false /* handleStrings */); err != nil {
				glog.Errorf("Unable to merge and upload CSV files for %s: %s", run, err)
			}
			// Cleanup created dir after the run completes.
			defer skutil.RemoveAll(filepath.Join(util.StorageDir, util.BenchmarkRunsDir, run))
		}
	}

	// Compare the resultant CSV files using csv_comparer.py
	noPatchCSVPath := filepath.Join(util.StorageDir, util.BenchmarkRunsDir, runIDNoPatch, runIDNoPatch+".output")
	withPatchCSVPath := filepath.Join(util.StorageDir, util.BenchmarkRunsDir, runIDWithPatch, runIDWithPatch+".output")
	htmlOutputDir := filepath.Join(util.StorageDir, util.ChromiumPerfRunsDir, *runID, "html")
	skutil.MkdirAll(htmlOutputDir, 0700)
	htmlRemoteDir := filepath.Join(remoteOutputDir, "html")
	htmlOutputLinkBase := util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, htmlRemoteDir) + "/"
	htmlOutputLink = htmlOutputLinkBase + "index.html"
	noPatchOutputLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, util.BenchmarkRunsDir, runIDNoPatch, "consolidated_outputs", runIDNoPatch+".output")
	withPatchOutputLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, util.BenchmarkRunsDir, runIDWithPatch, "consolidated_outputs", runIDWithPatch+".output")
	// Construct path to the csv_comparer python script.
	pathToCsvComparer := filepath.Join(pathToPyFiles, "csv_comparer.py")
	args := []string{
		pathToCsvComparer,
		"--csv_file1=" + noPatchCSVPath,
		"--csv_file2=" + withPatchCSVPath,
		"--output_html=" + htmlOutputDir,
		"--variance_threshold=" + strconv.FormatFloat(*varianceThreshold, 'f', 2, 64),
		"--discard_outliers=" + strconv.FormatFloat(*discardOutliers, 'f', 2, 64),
		"--absolute_url=" + htmlOutputLinkBase,
		"--requester_email=" + *emails,
		"--skia_patch_link=" + skiaPatchLink,
		"--chromium_patch_link=" + chromiumPatchLink,
		"--benchmark_patch_link=" + benchmarkPatchLink,
		"--description=" + *description,
		"--raw_csv_nopatch=" + noPatchOutputLink,
		"--raw_csv_withpatch=" + withPatchOutputLink,
		"--num_repeated=" + strconv.Itoa(*repeatBenchmark),
		"--target_platform=" + *targetPlatform,
		"--browser_args_nopatch=" + *browserExtraArgsNoPatch,
		"--browser_args_withpatch=" + *browserExtraArgsWithPatch,
		"--pageset_type=" + *pagesetType,
		"--chromium_hash=" + chromiumHash,
		"--skia_hash=" + skiaHash,
		"--missing_output_slaves=" + strings.Join(noOutputSlaves, " "),
		"--logs_link_prefix=" + fmt.Sprintf(util.SWARMING_RUN_ID_TASK_LINK_PREFIX_TEMPLATE, *runID, "chromium_perf_"),
	}
	err = util.ExecuteCmd("python", args, []string{}, util.CSV_COMPARER_TIMEOUT, nil, nil)
	if err != nil {
		glog.Errorf("Error running csv_comparer.py: %s", err)
		return
	}

	// Copy the HTML files to Google Storage.
	if err := gs.UploadDir(htmlOutputDir, htmlRemoteDir, true); err != nil {
		glog.Errorf("Could not upload %s to %s: %s", htmlOutputDir, htmlRemoteDir, err)
		return
	}

	taskCompletedSuccessfully = true
}

// getHashesFromBuild returns the Chromium and Skia hashes from a CT build string.
func getHashesFromBuild(chromiumBuild string) (string, string) {
	tokens := strings.Split(chromiumBuild, "-")
	return tokens[1], tokens[2]
}
<commit_msg>Handle stings in run_chromium_perf_on_workers<commit_after>// run_chromium_perf_on_workers is an application that runs the specified telemetry
// benchmark on all CT workers and uploads the results to Google Storage. The
// requester is emailed when the task is done.
package main

import (
	"database/sql"
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/skia-dev/glog"
	"go.skia.org/infra/ct/go/ctfe/chromium_perf"
	"go.skia.org/infra/ct/go/frontend"
	"go.skia.org/infra/ct/go/master_scripts/master_common"
	"go.skia.org/infra/ct/go/util"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/email"
	skutil "go.skia.org/infra/go/util"
)

const (
	MAX_PAGES_PER_SWARMING_BOT = 100
)

var (
	emails                    = flag.String("emails", "", "The comma separated email addresses to notify when the task is picked up and completes.")
	description               = flag.String("description", "", "The description of the run as entered by the requester.")
	gaeTaskID                 = flag.Int64("gae_task_id", -1, "The key of the App Engine task. This task will be updated when the task is completed.")
	pagesetType               = flag.String("pageset_type", "", "The type of pagesets to use. Eg: 10k, Mobile10k, All.")
	benchmarkName             = flag.String("benchmark_name", "", "The telemetry benchmark to run on the workers.")
	benchmarkExtraArgs        = flag.String("benchmark_extra_args", "", "The extra arguments that are passed to the specified benchmark.")
	browserExtraArgsNoPatch   = flag.String("browser_extra_args_nopatch", "", "The extra arguments that are passed to the browser while running the benchmark for the nopatch case.")
	browserExtraArgsWithPatch = flag.String("browser_extra_args_withpatch", "", "The extra arguments that are passed to the browser while running the benchmark for the withpatch case.")
	repeatBenchmark           = flag.Int("repeat_benchmark", 1, "The number of times the benchmark should be repeated. For skpicture_printer benchmark this value is always 1.")
	runInParallel             = flag.Bool("run_in_parallel", false, "Run the benchmark by bringing up multiple chrome instances in parallel.")
	targetPlatform            = flag.String("target_platform", util.PLATFORM_ANDROID, "The platform the benchmark will run on (Android / Linux).")
	runID                     = flag.String("run_id", "", "The unique run id (typically requester + timestamp).")
	varianceThreshold         = flag.Float64("variance_threshold", 0.0, "The variance threshold to use when comparing the resultant CSV files.")
	discardOutliers           = flag.Float64("discard_outliers", 0.0, "The percentage of outliers to discard when comparing the result CSV files.")

	taskCompletedSuccessfully = false

	htmlOutputLink      = util.MASTER_LOGSERVER_LINK
	skiaPatchLink       = util.MASTER_LOGSERVER_LINK
	chromiumPatchLink   = util.MASTER_LOGSERVER_LINK
	catapultPatchLink   = util.MASTER_LOGSERVER_LINK
	benchmarkPatchLink  = util.MASTER_LOGSERVER_LINK
	noPatchOutputLink   = util.MASTER_LOGSERVER_LINK
	withPatchOutputLink = util.MASTER_LOGSERVER_LINK
)

func sendEmail(recipients []string) {
	// Send completion email.
	emailSubject := fmt.Sprintf("Cluster telemetry chromium perf task has completed (%s)", *runID)
	failureHtml := ""
	viewActionMarkup := ""
	var err error

	if taskCompletedSuccessfully {
		if viewActionMarkup, err = email.GetViewActionMarkup(htmlOutputLink, "View Results", "Direct link to the HTML results"); err != nil {
			glog.Errorf("Failed to get view action markup: %s", err)
			return
		}
	} else {
		emailSubject += " with failures"
		failureHtml = util.GetFailureEmailHtml(*runID)
		if viewActionMarkup, err = email.GetViewActionMarkup(util.GetMasterLogLink(*runID), "View Failure", "Direct link to the master log"); err != nil {
			glog.Errorf("Failed to get view action markup: %s", err)
			return
		}
	}
	bodyTemplate := `
	The chromium perf %s benchmark task on %s pageset has completed.

	Run description: %s

	%s
	The HTML output with differences between the base run and the patch run is here.

	The patch(es) you specified are here:
	chromium/skia/catapult

	You can schedule more runs here.

	Thanks!
	`
	emailBody := fmt.Sprintf(bodyTemplate, *benchmarkName, *pagesetType, *description, failureHtml, htmlOutputLink, chromiumPatchLink, skiaPatchLink, catapultPatchLink, frontend.ChromiumPerfTasksWebapp)
	if err := util.SendEmailWithMarkup(recipients, emailSubject, emailBody, viewActionMarkup); err != nil {
		glog.Errorf("Error while sending email: %s", err)
		return
	}
}

func updateWebappTask() {
	vars := chromium_perf.UpdateVars{}
	vars.Id = *gaeTaskID
	vars.SetCompleted(taskCompletedSuccessfully)
	vars.Results = sql.NullString{String: htmlOutputLink, Valid: true}
	vars.NoPatchRawOutput = sql.NullString{String: noPatchOutputLink, Valid: true}
	vars.WithPatchRawOutput = sql.NullString{String: withPatchOutputLink, Valid: true}
	skutil.LogErr(frontend.UpdateWebappTaskV2(&vars))
}

func main() {
	defer common.LogPanic()
	master_common.Init()

	// Send start email.
	emailsArr := util.ParseEmails(*emails)
	emailsArr = append(emailsArr, util.CtAdmins...)
	if len(emailsArr) == 0 {
		glog.Error("At least one email address must be specified")
		return
	}
	skutil.LogErr(frontend.UpdateWebappTaskSetStarted(&chromium_perf.UpdateVars{}, *gaeTaskID))
	skutil.LogErr(util.SendTaskStartEmail(emailsArr, "Chromium perf", *runID, *description))
	// Ensure webapp is updated and email is sent even if task fails.
	defer updateWebappTask()
	defer sendEmail(emailsArr)
	// Cleanup dirs after run completes.
	defer skutil.RemoveAll(filepath.Join(util.StorageDir, util.ChromiumPerfRunsDir, *runID))
	// Finish with glog flush and how long the task took.
	defer util.TimeTrack(time.Now(), "Running chromium perf task on workers")
	defer glog.Flush()

	if *pagesetType == "" {
		glog.Error("Must specify --pageset_type")
		return
	}
	if *benchmarkName == "" {
		glog.Error("Must specify --benchmark_name")
		return
	}
	if *runID == "" {
		glog.Error("Must specify --run_id")
		return
	}
	if *description == "" {
		glog.Error("Must specify --description")
		return
	}

	// Instantiate GsUtil object.
	gs, err := util.NewGsUtil(nil)
	if err != nil {
		glog.Errorf("Could not instantiate gsutil object: %s", err)
		return
	}
	remoteOutputDir := filepath.Join(util.ChromiumPerfRunsDir, *runID)

	// Copy the patches to Google Storage.
	skiaPatchName := *runID + ".skia.patch"
	chromiumPatchName := *runID + ".chromium.patch"
	catapultPatchName := *runID + ".catapult.patch"
	benchmarkPatchName := *runID + ".benchmark.patch"
	for _, patchName := range []string{skiaPatchName, chromiumPatchName, catapultPatchName, benchmarkPatchName} {
		if err := gs.UploadFile(patchName, os.TempDir(), remoteOutputDir); err != nil {
			glog.Errorf("Could not upload %s to %s: %s", patchName, remoteOutputDir, err)
			return
		}
	}
	skiaPatchLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, remoteOutputDir, skiaPatchName)
	chromiumPatchLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, remoteOutputDir, chromiumPatchName)
	catapultPatchLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, remoteOutputDir, catapultPatchName)
	benchmarkPatchLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, remoteOutputDir, benchmarkPatchName)

	// Create the two required chromium builds (with patch and without the patch).
	chromiumBuilds, err := util.TriggerBuildRepoSwarmingTask(
		"build_chromium", *runID, "chromium", *targetPlatform, []string{},
		[]string{filepath.Join(remoteOutputDir, chromiumPatchName), filepath.Join(remoteOutputDir, skiaPatchName)},
		/*singlebuild*/ false, 3*time.Hour, 1*time.Hour)
	if err != nil {
		glog.Errorf("Error encountered when swarming build repo task: %s", err)
		return
	}
	if len(chromiumBuilds) != 2 {
		glog.Errorf("Expected 2 builds but instead got %d: %v.", len(chromiumBuilds), chromiumBuilds)
		return
	}
	chromiumBuildNoPatch := chromiumBuilds[0]
	chromiumBuildWithPatch := chromiumBuilds[1]

	// Parse out the Chromium and Skia hashes.
	chromiumHash, skiaHash := getHashesFromBuild(chromiumBuildNoPatch)

	// Archive, trigger and collect swarming tasks.
	isolateExtraArgs := map[string]string{
		"CHROMIUM_BUILD_NOPATCH":       chromiumBuildNoPatch,
		"CHROMIUM_BUILD_WITHPATCH":     chromiumBuildWithPatch,
		"RUN_ID":                       *runID,
		"BENCHMARK":                    *benchmarkName,
		"BENCHMARK_ARGS":               *benchmarkExtraArgs,
		"BROWSER_EXTRA_ARGS_NOPATCH":   *browserExtraArgsNoPatch,
		"BROWSER_EXTRA_ARGS_WITHPATCH": *browserExtraArgsWithPatch,
		"REPEAT_BENCHMARK":             strconv.Itoa(*repeatBenchmark),
		"RUN_IN_PARALLEL":              strconv.FormatBool(*runInParallel),
		"TARGET_PLATFORM":              *targetPlatform,
	}
	if err := util.TriggerSwarmingTask(*pagesetType, "chromium_perf", util.CHROMIUM_PERF_ISOLATE, *runID, 12*time.Hour, 1*time.Hour, util.USER_TASKS_PRIORITY, MAX_PAGES_PER_SWARMING_BOT, isolateExtraArgs, util.GOLO_WORKER_DIMENSIONS); err != nil {
		glog.Errorf("Error encountered when swarming tasks: %s", err)
		return
	}

	// If "--output-format=csv-pivot-table" was specified then merge all CSV files and upload.
	runIDNoPatch := fmt.Sprintf("%s-nopatch", *runID)
	runIDWithPatch := fmt.Sprintf("%s-withpatch", *runID)
	noOutputSlaves := []string{}
	pathToPyFiles := util.GetPathToPyFiles(false)
	for _, run := range []string{runIDNoPatch, runIDWithPatch} {
		if strings.Contains(*benchmarkExtraArgs, "--output-format=csv-pivot-table") {
			if noOutputSlaves, err = util.MergeUploadCSVFiles(run, pathToPyFiles, gs, util.PagesetTypeToInfo[*pagesetType].NumPages, MAX_PAGES_PER_SWARMING_BOT, true /* handleStrings */); err != nil {
				glog.Errorf("Unable to merge and upload CSV files for %s: %s", run, err)
			}
			// Cleanup created dir after the run completes.
			defer skutil.RemoveAll(filepath.Join(util.StorageDir, util.BenchmarkRunsDir, run))
		}
	}

	// Compare the resultant CSV files using csv_comparer.py
	noPatchCSVPath := filepath.Join(util.StorageDir, util.BenchmarkRunsDir, runIDNoPatch, runIDNoPatch+".output")
	withPatchCSVPath := filepath.Join(util.StorageDir, util.BenchmarkRunsDir, runIDWithPatch, runIDWithPatch+".output")
	htmlOutputDir := filepath.Join(util.StorageDir, util.ChromiumPerfRunsDir, *runID, "html")
	skutil.MkdirAll(htmlOutputDir, 0700)
	htmlRemoteDir := filepath.Join(remoteOutputDir, "html")
	htmlOutputLinkBase := util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, htmlRemoteDir) + "/"
	htmlOutputLink = htmlOutputLinkBase + "index.html"
	noPatchOutputLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, util.BenchmarkRunsDir, runIDNoPatch, "consolidated_outputs", runIDNoPatch+".output")
	withPatchOutputLink = util.GS_HTTP_LINK + filepath.Join(util.GSBucketName, util.BenchmarkRunsDir, runIDWithPatch, "consolidated_outputs", runIDWithPatch+".output")
	// Construct path to the csv_comparer python script.
	pathToCsvComparer := filepath.Join(pathToPyFiles, "csv_comparer.py")
	args := []string{
		pathToCsvComparer,
		"--csv_file1=" + noPatchCSVPath,
		"--csv_file2=" + withPatchCSVPath,
		"--output_html=" + htmlOutputDir,
		"--variance_threshold=" + strconv.FormatFloat(*varianceThreshold, 'f', 2, 64),
		"--discard_outliers=" + strconv.FormatFloat(*discardOutliers, 'f', 2, 64),
		"--absolute_url=" + htmlOutputLinkBase,
		"--requester_email=" + *emails,
		"--skia_patch_link=" + skiaPatchLink,
		"--chromium_patch_link=" + chromiumPatchLink,
		"--benchmark_patch_link=" + benchmarkPatchLink,
		"--description=" + *description,
		"--raw_csv_nopatch=" + noPatchOutputLink,
		"--raw_csv_withpatch=" + withPatchOutputLink,
		"--num_repeated=" + strconv.Itoa(*repeatBenchmark),
		"--target_platform=" + *targetPlatform,
		"--browser_args_nopatch=" + *browserExtraArgsNoPatch,
		"--browser_args_withpatch=" + *browserExtraArgsWithPatch,
		"--pageset_type=" + *pagesetType,
		"--chromium_hash=" + chromiumHash,
		"--skia_hash=" + skiaHash,
		"--missing_output_slaves=" + strings.Join(noOutputSlaves, " "),
		"--logs_link_prefix=" + fmt.Sprintf(util.SWARMING_RUN_ID_TASK_LINK_PREFIX_TEMPLATE, *runID, "chromium_perf_"),
	}
	err = util.ExecuteCmd("python", args, []string{}, util.CSV_COMPARER_TIMEOUT, nil, nil)
	if err != nil {
		glog.Errorf("Error running csv_comparer.py: %s", err)
		return
	}

	// Copy the HTML files to Google Storage.
	if err := gs.UploadDir(htmlOutputDir, htmlRemoteDir, true); err != nil {
		glog.Errorf("Could not upload %s to %s: %s", htmlOutputDir, htmlRemoteDir, err)
		return
	}

	taskCompletedSuccessfully = true
}

// getHashesFromBuild returns the Chromium and Skia hashes from a CT build string.
func getHashesFromBuild(chromiumBuild string) (string, string) {
	tokens := strings.Split(chromiumBuild, "-")
	return tokens[1], tokens[2]
}
<|endoftext|>


<commit_before>package flux

import (
	"log"
	"sync"
	"sync/atomic"
	"time"
)

//WaitInterface defines the flux.Wait interface method definitions
type WaitInterface interface {
	Add()
	Done()
	Count() int
	Flush()
	Then() ActionInterface
}

//ResetTimer runs a timer and performs an action
type ResetTimer struct {
	reset    chan struct{}
	kill     chan struct{}
	duration time.Duration
	do       *sync.Once
	init     func()
	done     func()
	state    int64
	started  int64
}

//NewResetTimer returns a new reset timer
func NewResetTimer(init func(), done func(), d time.Duration, run, boot bool) *ResetTimer {
	rs := &ResetTimer{
		reset:    make(chan struct{}),
		kill:     make(chan struct{}),
		duration: d,
		do:       new(sync.Once),
		init:     init,
		done:     done,
		state:    1,
		started:  0,
	}

	if run {
		rs.RunInit()
	}

	if boot {
		rs.handle()
	}

	return rs
}

//RunInit runs the init function
func (r *ResetTimer) RunInit() {
	if r.init != nil {
		r.init()
	}
}

//Add reset the timer threshold
func (r *ResetTimer) Add() {
	state := atomic.LoadInt64(&r.state)
	{
		log.Printf("Waiter checking State! State at %d", state)
		if state > 0 {
			r.reset <- struct{}{}
		} else {
			atomic.StoreInt64(&r.state, 1)
			{
				r.RunInit()
				r.handle()
			}
			atomic.StoreInt64(&r.started, 1)
		}
	}
}

//Close closes this timer
func (r *ResetTimer) Close() {
	state := atomic.LoadInt64(&r.started)
	{
		if state > 0 {
			close(r.kill)
			close(r.reset)
		}
	}
	atomic.StoreInt64(&r.state, 0)
	{
		r.kill = make(chan struct{})
		r.reset = make(chan struct{})
	}
	atomic.StoreInt64(&r.started, 0)
}

func (r *ResetTimer) makeTime() <-chan time.Time {
	return time.After(r.duration)
}

func (r *ResetTimer) handle() {
	go func() {
		threshold := r.makeTime()
	resetloop:
		for {
			select {
			case <-r.reset:
				log.Printf("ResetTimer: reseting timer by duration %+s", r.duration)
				threshold = r.makeTime()
			case <-threshold:
				atomic.StoreInt64(&r.state, 0)
				log.Printf("ResetTimer: timer expired with duration %+s", r.duration)
				r.done()
				break resetloop
			case <-r.kill:
				log.Printf("ResetTimer: timer killed using duration %+s", r.duration)
				atomic.StoreInt64(&r.state, 0)
				break resetloop
			}
		}

	}()
}

//SwitchInterface defines a flux.Switch interface method definition
type SwitchInterface interface {
	Switch()
	IsOn() bool
	WhenOn() ActionInterface
	WhenOff() ActionInterface
}

//WaitGen is a nice way of creating regenerative timers for use
//wait timers are once timers, once they are clocked out they are of no more use,to allow their nature which has its benefits we get to create WaitGen that generates a new once once a wait gen is over
type WaitGen struct {
	current WaitInterface
	gen     func() WaitInterface
}

//Make returns a new WaitInterface or returns the current once
func (w *WaitGen) Make() WaitInterface {
	if w.current != nil {
		return w.current
	}
	wt := w.gen()
	wt.Then().WhenOnly(func(_ interface{}) {
		w.current = nil
	})
	return wt
}

//NewTimeWaitGen returns a wait generator making a timewaiter
func NewTimeWaitGen(steps int, ms time.Duration, init func(WaitInterface)) *WaitGen {
	return &WaitGen{
		nil,
		func() WaitInterface {
			nt := NewTimeWait(steps, ms)
			init(nt)
			return nt
		},
	}
}

//NewSimpleWaitGen returns a wait generator making a timewaiter
func NewSimpleWaitGen(init func(WaitInterface)) *WaitGen {
	return &WaitGen{
		nil,
		func() WaitInterface {
			nt := NewWait()
			init(nt)
			return nt
		},
	}
}

//baseWait defines the base wait structure for all waiters
type baseWait struct {
	action ActionInterface
}

//Then returns an ActionInterface which gets fullfilled when this wait
//counter reaches zero
func (w *baseWait) Then() ActionInterface {
	return w.action.Wrap()
}

func newBaseWait() *baseWait {
	return &baseWait{NewAction()}
}

//TimeWait defines a time lock waiter
type TimeWait struct {
	*baseWait
	closer chan struct{}
	hits   int64
	max    int
	ms     time.Duration
	doonce *sync.Once
}

//NewTimeWait returns a new timer wait locker
//You specifiy two arguments:
//max int: the maximum number of time you want to check for idleness
//duration time.Duration: the time to check for each idle times and reduce
//until zero is reached then close
//eg. to do a 15seconds check for idleness
//NewTimeWait(15,time.Duration(1)*time.Second)
//eg. to do a 25 maximum check before closing per minute
//NewTimeWait(15,time.Duration(1)*time.Minute)
func NewTimeWait(max int, duration time.Duration) *TimeWait {

	tm := &TimeWait{
		newBaseWait(),
		make(chan struct{}),
		int64(max),
		max,
		duration,
		new(sync.Once),
	}

	// tm.Add()
	go tm.handle()

	return tm
}

//handle effects the necessary time process for checking and reducing the
//time checker for each duration of time,till the Waiter is done
func (w *TimeWait) handle() {
	var state int64
	atomic.StoreInt64(&state, 0)

	go func() {
		<-w.closer
		atomic.StoreInt64(&state, 1)
	}()

	for {
		time.Sleep(w.ms)

		bit := atomic.LoadInt64(&state)
		if bit > 0 {
			break
		}

		w.Done()
	}
}

//Flush drops the lock count and forces immediate unlocking of the wait
func (w *TimeWait) Flush() {
	w.doonce.Do(func() {
		close(w.closer)
		w.action.Fullfill(0)
		atomic.StoreInt64(&w.hits, 0)
	})
}

//Count returns the total left count to completed before unlock
func (w *TimeWait) Count() int {
	return int(atomic.LoadInt64(&w.hits))
}

//Add increments the lock state to the lock counter unless its already unlocked
func (w *TimeWait) Add() {
	if w.Count() < 0 || w.Count() >= w.max {
		return
	}

	atomic.AddInt64(&w.hits, 1)
}

//Done decrements the totalcount of this waitlocker by 1 until its below zero
//and fullfills with the 0 value
func (w *TimeWait) Done() {
	hits := atomic.LoadInt64(&w.hits)

	if hits < 0 {
		return
	}

	newhit := atomic.AddInt64(&w.hits, -1)
	log.Printf("TimeWait: Count Down now %d before %d!", newhit, hits)
	if int(newhit) <= 0 {
		w.Flush()
		log.Printf("TimeWait: Count Down Finished!")
	}
}

//Wait implements the WiatInterface for creating a wait lock which
//waits until the lock lockcount is finished then executes a action
//can only be used once, that is ,once the wait counter is -1,you cant add
//to it anymore
type Wait struct {
	*baseWait
	totalCount int64
}

//NewWait returns a new Wait instance for the WaitInterface
func NewWait() WaitInterface {
	return &Wait{newBaseWait(), int64(0)}
}

//Flush drops the lock count and forces immediate unlocking of the wait
func (w *Wait) Flush() {
	curr := int(atomic.LoadInt64(&w.totalCount))
	if curr < 0 {
		return
	}

	atomic.StoreInt64(&w.totalCount, 0)
	w.Done()
}

//Count returns the total left count to completed before unlock
func (w *Wait) Count() int {
	return int(atomic.LoadInt64(&w.totalCount))
}

//Add increments the lock state to the lock counter unless its already unlocked
func (w *Wait) Add() {
	curr := atomic.LoadInt64(&w.totalCount)

	if curr < 0 {
		return
	}

	atomic.AddInt64(&w.totalCount, 1)
}

//Done decrements the totalcount of this waitlocker by 1 until its below zero
//and fullfills with the 0 value
func (w *Wait) Done() {
	curr := atomic.LoadInt64(&w.totalCount)

	if curr < 0 {
		return
	}

	nc := atomic.AddInt64(&w.totalCount, -1)
	// log.Printf("Wait: Count Down now %d before %d", nc, curr)

	if int(nc) <= 0 {
		w.action.Fullfill(0)
	}
}
<commit_msg>fixing timer<commit_after>package flux

import (
	"log"
	"sync"
	"sync/atomic"
	"time"
)

//WaitInterface defines the flux.Wait interface method definitions
type WaitInterface interface {
	Add()
	Done()
	Count() int
	Flush()
	Then() ActionInterface
}

//ResetTimer runs a timer and performs an action
type ResetTimer struct {
	reset    chan struct{}
	kill     chan struct{}
	duration time.Duration
	do       *sync.Once
	init     func()
	done     func()
	state    int64
	started  int64
}

//NewResetTimer returns a new reset timer
func NewResetTimer(init func(), done func(), d time.Duration, run, boot bool) *ResetTimer {
	rs := &ResetTimer{
		reset:    make(chan struct{}),
		kill:     make(chan struct{}),
		duration: d,
		do:       new(sync.Once),
		init:     init,
		done:     done,
		state:    0,
		started:  1,
	}

	if run {
		rs.runInit()
	}

	if boot {
		rs.handle()
	}

	return rs
}

//RunInit runs the init function
func (r *ResetTimer) runInit() {
	if r.init != nil {
		r.init()
	}
}

//Add reset the timer threshold
func (r *ResetTimer) Add() {
	startd := int(atomic.LoadInt64(&r.started))
	state := int(atomic.LoadInt64(&r.state))

	if startd <= 0 {
		r.kill = make(chan struct{})
		r.reset = make(chan struct{})
	}

	log.Printf("Waiter checking State! State at %d Started %d", state, startd)

	if state > 0 {
		r.reset <- struct{}{}
	} else {
		r.runInit()
		r.handle()
		atomic.StoreInt64(&r.state, 1)
		atomic.StoreInt64(&r.started, 1)
	}
}

//Close closes this timer
func (r *ResetTimer) Close() {
	state := atomic.LoadInt64(&r.started)

	if state > 0 {
		close(r.kill)
		close(r.reset)
	}

	atomic.StoreInt64(&r.state, 0)
	atomic.StoreInt64(&r.started, 0)
}

func (r *ResetTimer) makeTime() <-chan time.Time {
	return time.After(r.duration)
}

func (r *ResetTimer) handle() {
	go func() {
		threshold := r.makeTime()
	resetloop:
		for {
			select {
			case <-r.reset:
				log.Printf("ResetTimer: reseting timer by duration %+s", r.duration)
				threshold = r.makeTime()
			case <-threshold:
				log.Printf("ResetTimer: timer expired with duration %+s", r.duration)
				r.done()
				atomic.StoreInt64(&r.state, 0)
				break resetloop
			case <-r.kill:
				log.Printf("ResetTimer: timer killed using duration %+s", r.duration)
				atomic.StoreInt64(&r.state, 0)
				break resetloop
			}
		}

	}()
}

//SwitchInterface defines a flux.Switch interface method definition
type SwitchInterface interface {
	Switch()
	IsOn() bool
	WhenOn() ActionInterface
	WhenOff() ActionInterface
}

//WaitGen is a nice way of creating regenerative timers for use
//wait timers are once timers, once they are clocked out they are of no more use,to allow their nature which has its benefits we get to create WaitGen that generates a new once once a wait gen is over
type WaitGen struct {
	current WaitInterface
	gen     func() WaitInterface
}

//Make returns a new WaitInterface or returns the current once
func (w *WaitGen) Make() WaitInterface {
	if w.current != nil {
		return w.current
	}
	wt := w.gen()
	wt.Then().WhenOnly(func(_ interface{}) {
		w.current = nil
	})
	return wt
}

//NewTimeWaitGen returns a wait generator making a timewaiter
func NewTimeWaitGen(steps int, ms time.Duration, init func(WaitInterface)) *WaitGen {
	return &WaitGen{
		nil,
		func() WaitInterface {
			nt := NewTimeWait(steps, ms)
			init(nt)
			return nt
		},
	}
}

//NewSimpleWaitGen returns a wait generator making a timewaiter
func NewSimpleWaitGen(init func(WaitInterface)) *WaitGen {
	return &WaitGen{
		nil,
		func() WaitInterface {
			nt := NewWait()
			init(nt)
			return nt
		},
	}
}

//baseWait defines the base wait structure for all waiters
type baseWait struct {
	action ActionInterface
}

//Then returns an ActionInterface which gets fullfilled when this wait
//counter reaches zero
func (w *baseWait) Then() ActionInterface {
	return w.action.Wrap()
}

func newBaseWait() *baseWait {
	return &baseWait{NewAction()}
}

//TimeWait defines a time lock waiter
type TimeWait struct {
	*baseWait
	closer chan struct{}
	hits   int64
	max    int
	ms     time.Duration
	doonce *sync.Once
}

//NewTimeWait returns a new timer wait locker
//You specifiy two arguments:
//max int: the maximum number of time you want to check for idleness
//duration time.Duration: the time to check for each idle times and reduce
//until zero is reached then close
//eg. to do a 15seconds check for idleness
//NewTimeWait(15,time.Duration(1)*time.Second)
//eg. to do a 25 maximum check before closing per minute
//NewTimeWait(15,time.Duration(1)*time.Minute)
func NewTimeWait(max int, duration time.Duration) *TimeWait {

	tm := &TimeWait{
		newBaseWait(),
		make(chan struct{}),
		int64(max),
		max,
		duration,
		new(sync.Once),
	}

	// tm.Add()
	go tm.handle()

	return tm
}

//handle effects the necessary time process for checking and reducing the
//time checker for each duration of time,till the Waiter is done
func (w *TimeWait) handle() {
	var state int64
	atomic.StoreInt64(&state, 0)

	go func() {
		<-w.closer
		atomic.StoreInt64(&state, 1)
	}()

	for {
		time.Sleep(w.ms)

		bit := atomic.LoadInt64(&state)
		if bit > 0 {
			break
		}

		w.Done()
	}
}

//Flush drops the lock count and forces immediate unlocking of the wait
func (w *TimeWait) Flush() {
	w.doonce.Do(func() {
		close(w.closer)
		w.action.Fullfill(0)
		atomic.StoreInt64(&w.hits, 0)
	})
}

//Count returns the total left count to completed before unlock
func (w *TimeWait) Count() int {
	return int(atomic.LoadInt64(&w.hits))
}

//Add increments the lock state to the lock counter unless its already unlocked
func (w *TimeWait) Add() {
	if w.Count() < 0 || w.Count() >= w.max {
		return
	}

	atomic.AddInt64(&w.hits, 1)
}

//Done decrements the totalcount of this waitlocker by 1 until its below zero
//and fullfills with the 0 value
func (w *TimeWait) Done() {
	hits := atomic.LoadInt64(&w.hits)

	if hits < 0 {
		return
	}

	newhit := atomic.AddInt64(&w.hits, -1)
	log.Printf("TimeWait: Count Down now %d before %d!", newhit, hits)
	if int(newhit) <= 0 {
		w.Flush()
		log.Printf("TimeWait: Count Down Finished!")
	}
}

//Wait implements the WiatInterface for creating a wait lock which
//waits until the lock lockcount is finished then executes a action
//can only be used once, that is ,once the wait counter is -1,you cant add
//to it anymore
type Wait struct {
	*baseWait
	totalCount int64
}

//NewWait returns a new Wait instance for the WaitInterface
func NewWait() WaitInterface {
	return &Wait{newBaseWait(), int64(0)}
}

//Flush drops the lock count and forces immediate unlocking of the wait
func (w *Wait) Flush() {
	curr := int(atomic.LoadInt64(&w.totalCount))
	if curr < 0 {
		return
	}

	atomic.StoreInt64(&w.totalCount, 0)
	w.Done()
}

//Count returns the total left count to completed before unlock
func (w *Wait) Count() int {
	return int(atomic.LoadInt64(&w.totalCount))
}

//Add increments the lock state to the lock counter unless its already unlocked
func (w *Wait) Add() {
	curr := atomic.LoadInt64(&w.totalCount)

	if curr < 0 {
		return
	}

	atomic.AddInt64(&w.totalCount, 1)
}

//Done decrements the totalcount of this waitlocker by 1 until its below zero
//and fullfills with the 0 value
func (w *Wait) Done() {
	curr := atomic.LoadInt64(&w.totalCount)

	if curr < 0 {
		return
	}

	nc := atomic.AddInt64(&w.totalCount, -1)
	// log.Printf("Wait: Count Down now %d before %d", nc, curr)

	if int(nc) <= 0 {
		w.action.Fullfill(0)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014 Mathias Monnerville. All rights reserved.
// Use of this source code is governed by a GPL
// license that can be found in the LICENSE file.

package mango

import (
	"encoding/json"
	"errors"
	"fmt"
)

// List of wallets.
type WalletList []*Wallet

// List of wallet's owners.
type ConsumerList []Consumer

// Money specifies which currency and amount (in cents!) to use in
// a payment transaction.
type Money struct {
	Currency string
	Amount   int // in cents, i.e 120 for 1.20 EUR
}

func (b Money) String() string {
	return fmt.Sprintf("%.2f %s", float64(b.Amount/100), b.Currency)
}

// Wallet stores all payins and tranfers from users in order to
// collect money.
type Wallet struct {
	ProcessIdent
	Owners      []string
	Description string
	Currency    string
	Balance     Money
	service     *MangoPay
}

func (u *Wallet) String() string {
	return struct2string(u)
}

// NewWallet creates a new wallet. Owners must have a well-defined Id. Empty Ids will
// return an error.
func (m *MangoPay) NewWallet(owners ConsumerList, desc string, currency string) (*Wallet, error) {
	all := []string{}
	for k, o := range owners {
		id := consumerId(o)
		if id == "" {
			return nil, errors.New(fmt.Sprintf("Empty Id for owner %d. Unable to create wallet.", k))
		}
		all = append(all, id)
	}
	w := &Wallet{
		Owners:      all,
		Description: desc,
		Currency:    currency,
	}
	w.service = m
	return w, nil
}

// Save creates or updates a legal user. The Create API is used
// if the user's Id is an empty string. The Edit API is used when
// the Id is a non-empty string.
func (w *Wallet) Save() error {
	var action mangoAction
	if w.Id == "" {
		action = actionCreateWallet
	} else {
		action = actionEditWallet
	}

	data := JsonObject{}
	j, err := json.Marshal(w)
	if err != nil {
		return err
	}
	if err := json.Unmarshal(j, &data); err != nil {
		return err
	}

	// Force float64 to int conversion after unmarshalling.
	for _, field := range []string{"CreationDate"} {
		data[field] = int(data[field].(float64))
	}

	// Fields not allowed when creating a wallet.
	if action == actionCreateWallet {
		delete(data, "Id")
	}
	delete(data, "CreationDate")
	delete(data, "Balance")

	if action == actionEditWallet {
		// Delete empty values so that existing ones don't get
		// overwritten with empty values.
		for k, v := range data {
			switch v.(type) {
			case string:
				if v.(string) == "" {
					delete(data, k)
				}
			case int:
				if v.(int) == 0 {
					delete(data, k)
				}
			}
		}
	}

	wallet, err := w.service.anyRequest(new(Wallet), action, data)
	if err != nil {
		return err
	}
	serv := w.service
	*w = *(wallet.(*Wallet))
	w.service = serv
	return nil
}

// Transactions returns a wallet's transactions.
func (w *Wallet) Transactions() (TransferList, error) {
	k, err := w.service.anyRequest(new(TransferList), actionFetchWalletTransactions, JsonObject{"Id": w.Id})
	if err != nil {
		return nil, err
	}
	return *(k.(*TransferList)), nil
}

// Wallet finds a legal user using the user_id attribute.
func (m *MangoPay) Wallet(id string) (*Wallet, error) {
	w, err := m.anyRequest(new(Wallet), actionFetchWallet, JsonObject{"Id": id})
	if err != nil {
		return nil, err
	}
	wallet := w.(*Wallet)
	wallet.service = m
	return wallet, nil
}

func (m *MangoPay) wallets(u Consumer) (WalletList, error) {
	id := consumerId(u)
	if id == "" {
		return nil, errors.New("user has empty Id")
	}
	trs, err := m.anyRequest(new(WalletList), actionFetchUserWallets, JsonObject{"Id": id})
	if err != nil {
		return nil, err
	}
	return *(trs.(*WalletList)), nil
}

// Wallet finds all user's wallets. Provided for convenience.
func (m *MangoPay) Wallets(user Consumer) (WalletList, error) {
	return m.wallets(user)
}
<commit_msg>Money formatting: float division fix<commit_after>// Copyright 2014 Mathias Monnerville. All rights reserved.
// Use of this source code is governed by a GPL
// license that can be found in the LICENSE file.

package mango

import (
	"encoding/json"
	"errors"
	"fmt"
)

// List of wallets.
type WalletList []*Wallet

// List of wallet's owners.
type ConsumerList []Consumer

// Money specifies which currency and amount (in cents!) to use in
// a payment transaction.
type Money struct {
	Currency string
	Amount   int // in cents, i.e 120 for 1.20 EUR
}

func (b Money) String() string {
	return fmt.Sprintf("%.2f %s", float64(b.Amount)/100, b.Currency)
}

// Wallet stores all payins and tranfers from users in order to
// collect money.
type Wallet struct {
	ProcessIdent
	Owners      []string
	Description string
	Currency    string
	Balance     Money
	service     *MangoPay
}

func (u *Wallet) String() string {
	return struct2string(u)
}

// NewWallet creates a new wallet. Owners must have a well-defined Id. Empty Ids will
// return an error.
func (m *MangoPay) NewWallet(owners ConsumerList, desc string, currency string) (*Wallet, error) {
	all := []string{}
	for k, o := range owners {
		id := consumerId(o)
		if id == "" {
			return nil, errors.New(fmt.Sprintf("Empty Id for owner %d. Unable to create wallet.", k))
		}
		all = append(all, id)
	}
	w := &Wallet{
		Owners:      all,
		Description: desc,
		Currency:    currency,
	}
	w.service = m
	return w, nil
}

// Save creates or updates a legal user. The Create API is used
// if the user's Id is an empty string. The Edit API is used when
// the Id is a non-empty string.
func (w *Wallet) Save() error {
	var action mangoAction
	if w.Id == "" {
		action = actionCreateWallet
	} else {
		action = actionEditWallet
	}

	data := JsonObject{}
	j, err := json.Marshal(w)
	if err != nil {
		return err
	}
	if err := json.Unmarshal(j, &data); err != nil {
		return err
	}

	// Force float64 to int conversion after unmarshalling.
	for _, field := range []string{"CreationDate"} {
		data[field] = int(data[field].(float64))
	}

	// Fields not allowed when creating a wallet.
	if action == actionCreateWallet {
		delete(data, "Id")
	}
	delete(data, "CreationDate")
	delete(data, "Balance")

	if action == actionEditWallet {
		// Delete empty values so that existing ones don't get
		// overwritten with empty values.
		for k, v := range data {
			switch v.(type) {
			case string:
				if v.(string) == "" {
					delete(data, k)
				}
			case int:
				if v.(int) == 0 {
					delete(data, k)
				}
			}
		}
	}

	wallet, err := w.service.anyRequest(new(Wallet), action, data)
	if err != nil {
		return err
	}
	serv := w.service
	*w = *(wallet.(*Wallet))
	w.service = serv
	return nil
}

// Transactions returns a wallet's transactions.
func (w *Wallet) Transactions() (TransferList, error) {
	k, err := w.service.anyRequest(new(TransferList), actionFetchWalletTransactions, JsonObject{"Id": w.Id})
	if err != nil {
		return nil, err
	}
	return *(k.(*TransferList)), nil
}

// Wallet finds a legal user using the user_id attribute.
func (m *MangoPay) Wallet(id string) (*Wallet, error) {
	w, err := m.anyRequest(new(Wallet), actionFetchWallet, JsonObject{"Id": id})
	if err != nil {
		return nil, err
	}
	wallet := w.(*Wallet)
	wallet.service = m
	return wallet, nil
}

func (m *MangoPay) wallets(u Consumer) (WalletList, error) {
	id := consumerId(u)
	if id == "" {
		return nil, errors.New("user has empty Id")
	}
	trs, err := m.anyRequest(new(WalletList), actionFetchUserWallets, JsonObject{"Id": id})
	if err != nil {
		return nil, err
	}
	return *(trs.(*WalletList)), nil
}

// Wallet finds all user's wallets. Provided for convenience.
func (m *MangoPay) Wallets(user Consumer) (WalletList, error) {
	return m.wallets(user)
}
<|endoftext|>


<commit_before>package transport

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestValidateSecureEndpoints(t *testing.T) {
	tlsInfo, err := createSelfCert(t)
	if err != nil {
		t.Fatalf("unable to create cert: %v", err)
	}

	remoteAddr := func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(r.RemoteAddr))
	}
	srv := httptest.NewServer(http.HandlerFunc(remoteAddr))
	defer srv.Close()

	insecureEps := []string{
		"http://" + srv.Listener.Addr().String(),
		"invalid remote address",
	}
	if _, err := ValidateSecureEndpoints(*tlsInfo, insecureEps); err == nil || !strings.Contains(err.Error(), "is insecure") {
		t.Error("validate secure endpoints should fail")
	}

	secureEps := []string{
		"https://" + srv.Listener.Addr().String(),
	}
	if _, err := ValidateSecureEndpoints(*tlsInfo, secureEps); err != nil {
		t.Error("validate secure endpoints should succeed")
	}
}
<commit_msg>Update client/pkg/transport/tls_test.go<commit_after>// Copyright 2022 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package transport

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestValidateSecureEndpoints(t *testing.T) {
	tlsInfo, err := createSelfCert(t)
	if err != nil {
		t.Fatalf("unable to create cert: %v", err)
	}

	remoteAddr := func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(r.RemoteAddr))
	}
	srv := httptest.NewServer(http.HandlerFunc(remoteAddr))
	defer srv.Close()

	insecureEps := []string{
		"http://" + srv.Listener.Addr().String(),
		"invalid remote address",
	}
	if _, err := ValidateSecureEndpoints(*tlsInfo, insecureEps); err == nil || !strings.Contains(err.Error(), "is insecure") {
		t.Error("validate secure endpoints should fail")
	}

	secureEps := []string{
		"https://" + srv.Listener.Addr().String(),
	}
	if _, err := ValidateSecureEndpoints(*tlsInfo, secureEps); err != nil {
		t.Error("validate secure endpoints should succeed")
	}
}
<|endoftext|>


<commit_before>package wechat

import (
	"net/http"
	"sync"

	"github.com/silenceper/wechat/cache"
	"github.com/silenceper/wechat/context"
	"github.com/silenceper/wechat/js"
	"github.com/silenceper/wechat/material"
	"github.com/silenceper/wechat/menu"
	"github.com/silenceper/wechat/oauth"
	"github.com/silenceper/wechat/server"
	"github.com/silenceper/wechat/template"
	"github.com/silenceper/wechat/user"
	"github.com/silenceper/wechat/pay"
)

// Wechat struct
type Wechat struct {
	Context *context.Context
}

// Config for user
type Config struct {
	AppID          string
	AppSecret      string
	Token          string
	EncodingAESKey string
	PayMchID       string  //支付 - 商户 ID
	PayNotifyURL   string  //支付 - 接受微信支付结果通知的接口地址
	PayKey         string  //支付 - 商户后台设置的支付 key
	Cache          cache.Cache
}

// NewWechat init
func NewWechat(cfg *Config) *Wechat {
	context := new(context.Context)
	copyConfigToContext(cfg, context)
	return &Wechat{context}
}

func copyConfigToContext(cfg *Config, context *context.Context) {
	context.AppID = cfg.AppID
	context.AppSecret = cfg.AppSecret
	context.Token = cfg.Token
	context.EncodingAESKey = cfg.EncodingAESKey
	context.Cache = cfg.Cache
	context.SetAccessTokenLock(new(sync.RWMutex))
	context.SetJsAPITicketLock(new(sync.RWMutex))
}

// GetServer 消息管理
func (wc *Wechat) GetServer(req *http.Request, writer http.ResponseWriter) *server.Server {
	wc.Context.Request = req
	wc.Context.Writer = writer
	return server.NewServer(wc.Context)
}

//GetAccessToken 获取access_token
func (wc *Wechat) GetAccessToken() (string, error) {
	return wc.Context.GetAccessToken()
}

// GetOauth oauth2网页授权
func (wc *Wechat) GetOauth() *oauth.Oauth {
	return oauth.NewOauth(wc.Context)
}

// GetMaterial 素材管理
func (wc *Wechat) GetMaterial() *material.Material {
	return material.NewMaterial(wc.Context)
}

// GetJs js-sdk配置
func (wc *Wechat) GetJs() *js.Js {
	return js.NewJs(wc.Context)
}

// GetMenu 菜单管理接口
func (wc *Wechat) GetMenu() *menu.Menu {
	return menu.NewMenu(wc.Context)
}

// GetUser 用户管理接口
func (wc *Wechat) GetUser() *user.User {
	return user.NewUser(wc.Context)
}

// GetTemplate 模板消息接口
func (wc *Wechat) GetTemplate() *template.Template {
	return template.NewTemplate(wc.Context)
}

// GetPay 返回支付消息的实例
func (wc *Wechat) GetPay() *pay.Pay {
	return pay.NewPay(wc.Context)
}<commit_msg>add config copy<commit_after>package wechat

import (
	"net/http"
	"sync"

	"github.com/silenceper/wechat/cache"
	"github.com/silenceper/wechat/context"
	"github.com/silenceper/wechat/js"
	"github.com/silenceper/wechat/material"
	"github.com/silenceper/wechat/menu"
	"github.com/silenceper/wechat/oauth"
	"github.com/silenceper/wechat/server"
	"github.com/silenceper/wechat/template"
	"github.com/silenceper/wechat/user"
	"github.com/silenceper/wechat/pay"
)

// Wechat struct
type Wechat struct {
	Context *context.Context
}

// Config for user
type Config struct {
	AppID          string
	AppSecret      string
	Token          string
	EncodingAESKey string
	PayMchID       string  //支付 - 商户 ID
	PayNotifyURL   string  //支付 - 接受微信支付结果通知的接口地址
	PayKey         string  //支付 - 商户后台设置的支付 key
	Cache          cache.Cache
}

// NewWechat init
func NewWechat(cfg *Config) *Wechat {
	context := new(context.Context)
	copyConfigToContext(cfg, context)
	return &Wechat{context}
}

func copyConfigToContext(cfg *Config, context *context.Context) {
	context.AppID = cfg.AppID
	context.AppSecret = cfg.AppSecret
	context.Token = cfg.Token
	context.EncodingAESKey = cfg.EncodingAESKey
	context.PayMchID = cfg.PayMchID
	context.PayKey = cfg.PayKey
	context.PayNotifyURL = cfg.PayNotifyURL
	context.Cache = cfg.Cache
	context.SetAccessTokenLock(new(sync.RWMutex))
	context.SetJsAPITicketLock(new(sync.RWMutex))
}

// GetServer 消息管理
func (wc *Wechat) GetServer(req *http.Request, writer http.ResponseWriter) *server.Server {
	wc.Context.Request = req
	wc.Context.Writer = writer
	return server.NewServer(wc.Context)
}

//GetAccessToken 获取access_token
func (wc *Wechat) GetAccessToken() (string, error) {
	return wc.Context.GetAccessToken()
}

// GetOauth oauth2网页授权
func (wc *Wechat) GetOauth() *oauth.Oauth {
	return oauth.NewOauth(wc.Context)
}

// GetMaterial 素材管理
func (wc *Wechat) GetMaterial() *material.Material {
	return material.NewMaterial(wc.Context)
}

// GetJs js-sdk配置
func (wc *Wechat) GetJs() *js.Js {
	return js.NewJs(wc.Context)
}

// GetMenu 菜单管理接口
func (wc *Wechat) GetMenu() *menu.Menu {
	return menu.NewMenu(wc.Context)
}

// GetUser 用户管理接口
func (wc *Wechat) GetUser() *user.User {
	return user.NewUser(wc.Context)
}

// GetTemplate 模板消息接口
func (wc *Wechat) GetTemplate() *template.Template {
	return template.NewTemplate(wc.Context)
}

// GetPay 返回支付消息的实例
func (wc *Wechat) GetPay() *pay.Pay {
	return pay.NewPay(wc.Context)
}<|endoftext|>


<commit_before>package extra

import (
	seabird "github.com/belak/go-seabird"
)

func init() {
	seabird.RegisterPlugin("mentions", newMentionsPlugin)
}

func newMentionsPlugin(b *seabird.Bot) error {
	mm := b.MentionMux()

	mm.Event(mentionsCallback)

	return nil
}

func mentionsCallback(r *seabird.Request) {
	switch r.Message.Trailing() {
	case "ping":
		r.MentionReply("pong")
	case "scoobysnack", "scooby snack":
		r.Reply("Scooby Dooby Doo!")
	case "botsnack", "bot snack":
		r.Reply(":)")
	}
}
<commit_msg>mention: add pizzahousesnack<commit_after>package extra

import (
	seabird "github.com/belak/go-seabird"
)

func init() {
	seabird.RegisterPlugin("mentions", newMentionsPlugin)
}

func newMentionsPlugin(b *seabird.Bot) error {
	mm := b.MentionMux()

	mm.Event(mentionsCallback)

	return nil
}

func mentionsCallback(r *seabird.Request) {
	switch r.Message.Trailing() {
	case "ping":
		r.MentionReply("pong")
	case "scoobysnack", "scooby snack":
		r.Reply("Scooby Dooby Doo!")
	case "botsnack", "bot snack":
		r.Reply(":)")
	case "pizzahousesnack":
		r.Reply("HECK YEAHHHHHHHHHHHH OMG I LOVE U THE WORLD IS GREAT")
	}
}
<|endoftext|>


<commit_before>package disruptor

// Cursors should be a party of the same backing array to keep them as close together as possible:
// https://news.ycombinator.com/item?id=7800825
type (
	Wireup struct {
		capacity int64
		groups   [][]Consumer
		cursors  []*Cursor // backing array keeps cursors (with padding) in contiguous memory
	}
)

func Configure(capacity int64) Wireup {
	return Wireup{
		capacity: capacity,
		groups:   [][]Consumer{},
		cursors:  []*Cursor{NewCursor()},
	}
}

func (this Wireup) WithConsumerGroup(consumers ...Consumer) Wireup {
	if len(consumers) == 0 {
		return this
	}

	target := make([]Consumer, len(consumers))
	copy(target, consumers)

	for i := 0; i < len(consumers); i++ {
		this.cursors = append(this.cursors, NewCursor())
	}

	this.groups = append(this.groups, target)
	return this
}

func (this Wireup) Build() Disruptor {
	allReaders := []*Reader{}
	written := this.cursors[0]
	var upstream Barrier = this.cursors[0]
	cursorIndex := 1 // 0 index is reserved for the writer Cursor

	for groupIndex, group := range this.groups {
		groupReaders, groupBarrier := this.buildReaders(groupIndex, cursorIndex, written, upstream)
		for _, item := range groupReaders {
			allReaders = append(allReaders, item)
		}
		upstream = groupBarrier
		cursorIndex += len(group)
	}

	writer := NewWriter(written, upstream, this.capacity)
	return Disruptor{writer: writer, readers: allReaders}
}

func (this Wireup) BuildShared() SharedDisruptor {
	allReaders := []*Reader{}
	written := this.cursors[0]
	writerBarrier := NewSharedWriterBarrier(written, this.capacity)
	var upstream Barrier = writerBarrier
	cursorIndex := 1 // 0 index is reserved for the writer Cursor

	for groupIndex, group := range this.groups {
		groupReaders, groupBarrier := this.buildReaders(groupIndex, cursorIndex, written, upstream)
		for _, item := range groupReaders {
			allReaders = append(allReaders, item)
		}
		upstream = groupBarrier
		cursorIndex += len(group)
	}

	writer := NewSharedWriter(writerBarrier, upstream)
	return SharedDisruptor{writer: writer, readers: allReaders}
}

func (this Wireup) buildReaders(consumerIndex, cursorIndex int, written *Cursor, upstream Barrier) ([]*Reader, Barrier) {
	barrierCursors := []*Cursor{}
	readers := []*Reader{}

	for _, consumer := range this.groups[consumerIndex] {
		cursor := this.cursors[cursorIndex]
		barrierCursors = append(barrierCursors, cursor)
		reader := NewReader(cursor, written, upstream, consumer)
		readers = append(readers, reader)
		cursorIndex++
	}

	if len(this.groups[consumerIndex]) == 1 {
		return readers, barrierCursors[0]
	} else {
		return readers, NewCompositeBarrier(barrierCursors...)
	}
}
<commit_msg>Variadic function.<commit_after>package disruptor

// Cursors should be a party of the same backing array to keep them as close together as possible:
// https://news.ycombinator.com/item?id=7800825
type (
	Wireup struct {
		capacity int64
		groups   [][]Consumer
		cursors  []*Cursor // backing array keeps cursors (with padding) in contiguous memory
	}
)

func Configure(capacity int64) Wireup {
	return Wireup{
		capacity: capacity,
		groups:   [][]Consumer{},
		cursors:  []*Cursor{NewCursor()},
	}
}

func (this Wireup) WithConsumerGroup(consumers ...Consumer) Wireup {
	if len(consumers) == 0 {
		return this
	}

	target := make([]Consumer, len(consumers))
	copy(target, consumers)

	for i := 0; i < len(consumers); i++ {
		this.cursors = append(this.cursors, NewCursor())
	}

	this.groups = append(this.groups, target)
	return this
}

func (this Wireup) Build() Disruptor {
	allReaders := []*Reader{}
	written := this.cursors[0]
	var upstream Barrier = this.cursors[0]
	cursorIndex := 1 // 0 index is reserved for the writer Cursor

	for groupIndex, group := range this.groups {
		groupReaders, groupBarrier := this.buildReaders(groupIndex, cursorIndex, written, upstream)
		allReaders = append(allReaders, groupReaders...)
		upstream = groupBarrier
		cursorIndex += len(group)
	}

	writer := NewWriter(written, upstream, this.capacity)
	return Disruptor{writer: writer, readers: allReaders}
}

func (this Wireup) BuildShared() SharedDisruptor {
	allReaders := []*Reader{}
	written := this.cursors[0]
	writerBarrier := NewSharedWriterBarrier(written, this.capacity)
	var upstream Barrier = writerBarrier
	cursorIndex := 1 // 0 index is reserved for the writer Cursor

	for groupIndex, group := range this.groups {
		groupReaders, groupBarrier := this.buildReaders(groupIndex, cursorIndex, written, upstream)
		for _, item := range groupReaders {
			allReaders = append(allReaders, item)
		}
		upstream = groupBarrier
		cursorIndex += len(group)
	}

	writer := NewSharedWriter(writerBarrier, upstream)
	return SharedDisruptor{writer: writer, readers: allReaders}
}

func (this Wireup) buildReaders(consumerIndex, cursorIndex int, written *Cursor, upstream Barrier) ([]*Reader, Barrier) {
	barrierCursors := []*Cursor{}
	readers := []*Reader{}

	for _, consumer := range this.groups[consumerIndex] {
		cursor := this.cursors[cursorIndex]
		barrierCursors = append(barrierCursors, cursor)
		reader := NewReader(cursor, written, upstream, consumer)
		readers = append(readers, reader)
		cursorIndex++
	}

	if len(this.groups[consumerIndex]) == 1 {
		return readers, barrierCursors[0]
	} else {
		return readers, NewCompositeBarrier(barrierCursors...)
	}
}
<|endoftext|>


<commit_before>package gokiq

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"syscall"
	"time"

	"github.com/garyburd/redigo/redis"
)

type Job struct {
	Type  string        `json:"class"`
	Args  []interface{} `json:"args"`
	Queue string        `json:"queue"`
	ID    string        `json:"jid"`

	Retry      interface{} `json:"retry"` // can be int (number of retries) or bool (true means default)
	MaxRetries int
	RetryCount int `json:"retry_count,omitmepty"`
	RetriedAt  int `json:"retried_at,omitempty"`

	ErrorMessage string `json:"error_message,omitempty"`
	ErrorType    string `json:"error_class,omitempty"`
	FailedAt     string `json:"failed_at,omitempty"`
}

type message struct {
	job *Job
	die bool
}

const (
	TimestampFormat = "2006-01-02 15:04:05 MST"
	redisTimeout    = 1
	maxIdleRedis    = 1
)

type QueueConfig map[string]int

type Worker interface {
	Perform([]interface{}) error
}

var Workers = NewWorkerConfig()

type WorkerConfig struct {
	RedisServer string
	Queues      QueueConfig
	WorkerCount int
	ReportError func(error, *Job)

	workerMapping map[string]reflect.Type
	randomQueues  []string
	redisPool     *redis.Pool
	workQueue     chan message
	doneQueue     chan bool
	done          bool
	sync.Mutex
}

func NewWorkerConfig() *WorkerConfig {
	return &WorkerConfig{
		ReportError:   func(error, *Job) {},
		workerMapping: make(map[string]reflect.Type),
		workQueue:     make(chan message),
		doneQueue:     make(chan bool),
	}
}

func (w *WorkerConfig) Register(name string, worker Worker) {
	w.workerMapping[name] = reflect.Indirect(reflect.ValueOf(worker)).Type()
}

func (w *WorkerConfig) Run() {
	w.denormalizeQueues()
	w.connectRedis()

	for i := 0; i < w.WorkerCount; i++ {
		go w.worker()
	}

	go w.retryScheduler()
	go w.quitHandler()

	for {
		if w.done {
			w.Lock() // we're done, so block until quitHandler() calls os.Exit()
		}

		w.Lock() // don't let quitHandler() stop us in the middle of the iteration
		msg, err := redis.Bytes(w.redisQuery("BLPOP", append(w.queueList(), redisTimeout)))
		if err != nil {
			w.handleError(err)
			time.Sleep(redisTimeout * time.Second) // likely a transient redis error, sleep before retrying
			continue
		}

		job := &Job{}
		json.Unmarshal(msg, job)
		if err != nil {
			w.handleError(err)
			continue
		}

		w.workQueue <- message{job: job}
		w.Unlock()
	}
}

// create a slice of queues with duplicates using the assigned frequencies
func (w *WorkerConfig) denormalizeQueues() {
	for queue, x := range w.Queues {
		for i := 0; i < x; i++ {
			w.randomQueues = append(w.randomQueues, "queue:"+queue)
		}
	}
}

// get a random slice of unique queues from the slice of denormalized queues
func (w *WorkerConfig) queueList() []interface{} {
	size := len(w.Queues)
	res := make([]interface{}, 0, size)
	queues := make(map[string]bool, size)

	indices := rand.Perm(len(w.randomQueues))[:size]
	for _, i := range indices {
		queue := w.randomQueues[i]
		if _, ok := queues[queue]; !ok {
			queues[queue] = true
			res = append(res, queue)
		}
	}

	return res
}

func (w *WorkerConfig) handleError(err error) {
	// TODO: log message to stdout
	w.ReportError(err, nil)
}

// checks the sorted set of scheduled retries and queues them when it's time
func (w *WorkerConfig) retryScheduler() {
	for _ = range time.Tick(time.Second) {

	}
}

// listens for SIGINT, SIGTERM, and SIGQUIT to perform a clean shutdown
func (w *WorkerConfig) quitHandler() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	signal.Notify(c, syscall.SIGQUIT)

	for _ = range c {
		w.Lock()           // wait for the current runner iteration to finish
		w.done = true      // tell the runner that we're done
		close(w.workQueue) // tell worker goroutines to stop after they finish their current job
		for i := 0; i < w.WorkerCount; i++ {
			<-w.doneQueue // wait for workers to finish
		}
		os.Exit(0)
	}
}

func (w *WorkerConfig) connectRedis() {
	w.redisPool = redis.NewPool(func() (redis.Conn, error) {
		return redis.Dial("tcp", w.RedisServer)
	}, maxIdleRedis)
}

func (w *WorkerConfig) redisQuery(command string, args ...interface{}) (interface{}, error) {
	conn := w.redisPool.Get()
	defer conn.Close()
	return conn.Do(command, args...)
}

func (w *WorkerConfig) worker() {
	for msg := range w.workQueue {
		if msg.die {
			return
		}

		job := msg.job
		typ, ok := w.workerMapping[msg.job.Type]
		if !ok {
			err := fmt.Errorf("Unknown worker type: %s", job.Type)
			w.scheduleRetry(job, err)
			continue
		}

		// wrap Perform() in a function so that we can recover from panics
		var err error
		func() {
			defer func() {
				if r := recover(); r != nil {
					// TODO: log stack trace
					err = panicToError(r)
				}
			}()
			err = reflect.New(typ).Interface().(Worker).Perform(msg.job.Args)
		}()
		if err != nil {
			w.scheduleRetry(job, err)
		}
	}
	w.doneQueue <- true
}

func (w *WorkerConfig) scheduleRetry(job *Job, err error) {
	w.ReportError(err, job)
	// set failure details
	// determine next retry and add to retry set
	// log failure stat
}

func panicToError(err interface{}) error {
	if str, ok := err.(string); ok {
		return fmt.Errorf(str)
	}
	return err.(error)
}
<commit_msg>Fix up locking semantics<commit_after>package gokiq

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"syscall"
	"time"

	"github.com/garyburd/redigo/redis"
)

type Job struct {
	Type  string        `json:"class"`
	Args  []interface{} `json:"args"`
	Queue string        `json:"queue"`
	ID    string        `json:"jid"`

	Retry      interface{} `json:"retry"` // can be int (number of retries) or bool (true means default)
	MaxRetries int
	RetryCount int `json:"retry_count,omitmepty"`
	RetriedAt  int `json:"retried_at,omitempty"`

	ErrorMessage string `json:"error_message,omitempty"`
	ErrorType    string `json:"error_class,omitempty"`
	FailedAt     string `json:"failed_at,omitempty"`
}

type message struct {
	job *Job
	die bool
}

const (
	TimestampFormat = "2006-01-02 15:04:05 MST"
	redisTimeout    = 1
	maxIdleRedis    = 1
)

type QueueConfig map[string]int

type Worker interface {
	Perform([]interface{}) error
}

var Workers = NewWorkerConfig()

type WorkerConfig struct {
	RedisServer string
	Queues      QueueConfig
	WorkerCount int
	ReportError func(error, *Job)

	workerMapping map[string]reflect.Type
	randomQueues  []string
	redisPool     *redis.Pool
	workQueue     chan message
	doneQueue     chan bool
	sync.Mutex
}

func NewWorkerConfig() *WorkerConfig {
	return &WorkerConfig{
		ReportError:   func(error, *Job) {},
		workerMapping: make(map[string]reflect.Type),
		workQueue:     make(chan message),
		doneQueue:     make(chan bool),
	}
}

func (w *WorkerConfig) Register(name string, worker Worker) {
	w.workerMapping[name] = reflect.Indirect(reflect.ValueOf(worker)).Type()
}

func (w *WorkerConfig) Run() {
	w.denormalizeQueues()
	w.connectRedis()

	for i := 0; i < w.WorkerCount; i++ {
		go w.worker()
	}

	go w.retryScheduler()
	go w.quitHandler()

	for {
		w.Lock() // don't let quitHandler() stop us in the middle of a job
		msg, err := redis.Bytes(w.redisQuery("BLPOP", append(w.queueList(), redisTimeout)))
		if err != nil {
			w.handleError(err)
			time.Sleep(redisTimeout * time.Second) // likely a transient redis error, sleep before retrying
			continue
		}

		job := &Job{}
		json.Unmarshal(msg, job)
		if err != nil {
			w.handleError(err)
			continue
		}
		w.workQueue <- message{job: job}

		w.Unlock()
		time.Sleep(time.Microsecond) // give quitHandler() time to grab the lock
	}
}

// create a slice of queues with duplicates using the assigned frequencies
func (w *WorkerConfig) denormalizeQueues() {
	for queue, x := range w.Queues {
		for i := 0; i < x; i++ {
			w.randomQueues = append(w.randomQueues, "queue:"+queue)
		}
	}
}

// get a random slice of unique queues from the slice of denormalized queues
func (w *WorkerConfig) queueList() []interface{} {
	size := len(w.Queues)
	res := make([]interface{}, 0, size)
	queues := make(map[string]bool, size)

	indices := rand.Perm(len(w.randomQueues))[:size]
	for _, i := range indices {
		queue := w.randomQueues[i]
		if _, ok := queues[queue]; !ok {
			queues[queue] = true
			res = append(res, queue)
		}
	}

	return res
}

func (w *WorkerConfig) handleError(err error) {
	// TODO: log message to stdout
	w.ReportError(err, nil)
}

// checks the sorted set of scheduled retries and queues them when it's time
func (w *WorkerConfig) retryScheduler() {
	for _ = range time.Tick(time.Second) {

	}
}

// listens for SIGINT, SIGTERM, and SIGQUIT to perform a clean shutdown
func (w *WorkerConfig) quitHandler() {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	signal.Notify(c, syscall.SIGQUIT)

	for _ = range c {
		w.Lock()           // wait for the current runner iteration to finish and stop it from continuing
		close(w.workQueue) // tell worker goroutines to stop after they finish their current job
		for i := 0; i < w.WorkerCount; i++ {
			<-w.doneQueue // wait for workers to finish
		}
		os.Exit(0)
	}
}

func (w *WorkerConfig) connectRedis() {
	w.redisPool = redis.NewPool(func() (redis.Conn, error) {
		return redis.Dial("tcp", w.RedisServer)
	}, maxIdleRedis)
}

func (w *WorkerConfig) redisQuery(command string, args ...interface{}) (interface{}, error) {
	conn := w.redisPool.Get()
	defer conn.Close()
	return conn.Do(command, args...)
}

func (w *WorkerConfig) worker() {
	for msg := range w.workQueue {
		if msg.die {
			return
		}

		job := msg.job
		typ, ok := w.workerMapping[msg.job.Type]
		if !ok {
			err := fmt.Errorf("Unknown worker type: %s", job.Type)
			w.scheduleRetry(job, err)
			continue
		}

		// wrap Perform() in a function so that we can recover from panics
		var err error
		func() {
			defer func() {
				if r := recover(); r != nil {
					// TODO: log stack trace
					err = panicToError(r)
				}
			}()
			err = reflect.New(typ).Interface().(Worker).Perform(msg.job.Args)
		}()
		if err != nil {
			w.scheduleRetry(job, err)
		}
	}
	w.doneQueue <- true
}

func (w *WorkerConfig) scheduleRetry(job *Job, err error) {
	w.ReportError(err, job)
	// set failure details
	// determine next retry and add to retry set
	// log failure stat
}

func panicToError(err interface{}) error {
	if str, ok := err.(string); ok {
		return fmt.Errorf(str)
	}
	return err.(error)
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package amass

import (
	"bufio"
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/gamexg/proxyclient"
)

// AmassConfig - Passes along optional configurations
type AmassConfig struct {
	sync.Mutex

	// The channel that will receive the results
	Output chan *AmassRequest

	// The ASNs that the enumeration will target
	ASNs []int

	// The CIDRs that the enumeration will target
	CIDRs []*net.IPNet

	// The IPs that the enumeration will target
	IPs []net.IP

	// The ports that will be checked for certificates
	Ports []int

	// The list of words to use when generating names
	Wordlist []string

	// Will the enumeration including brute forcing techniques
	BruteForcing bool

	// Will recursive brute forcing be performed?
	Recursive bool

	// Minimum number of subdomain discoveries before performing recursive brute forcing
	MinForRecursive int

	// Will discovered subdomain name alterations be generated?
	Alterations bool

	// A blacklist of subdomain names that will not be investigated
	Blacklist []string

	// Sets the maximum number of DNS queries per minute
	Frequency time.Duration

	// Preferred DNS resolvers identified by the user
	Resolvers []string

	// Indicate that Amass cannot add domains to the config
	AdditionalDomains bool

	// The root domain names that the enumeration will target
	domains []string

	// Is responsible for performing simple DNS resolutions
	dns *queries

	// Handles selecting the next DNS resolver to be used
	resolver *resolvers

	// Performs lookups of root domain names from subdomain names
	domainLookup *DomainLookup

	// Detects DNS wildcards
	wildcards *Wildcards

	// The optional proxy connection for the enumeration to use
	proxy proxyclient.ProxyClient
}

func (c *AmassConfig) Setup() {
	// Setup the services potentially needed by all of amass
	c.dns = newQueriesSubsystem(c)
	c.domainLookup = NewDomainLookup(c)
	c.wildcards = NewWildcardDetection(c)
	c.resolver = newResolversSubsystem(c)
}

func (c *AmassConfig) AddDomains(names []string) {
	c.Lock()
	defer c.Unlock()

	c.domains = UniqueAppend(c.domains, names...)
}

func (c *AmassConfig) Domains() []string {
	c.Lock()
	defer c.Unlock()

	return c.domains
}

func (c *AmassConfig) Blacklisted(name string) bool {
	var resp bool

	for _, bl := range c.Blacklist {
		if match := strings.HasSuffix(name, bl); match {
			resp = true
			break
		}
	}
	return resp
}

func CheckConfig(config *AmassConfig) error {
	if len(config.Wordlist) == 0 {
		return errors.New("The configuration contains no wordlist")
	}

	if config.Frequency < DefaultConfig().Frequency {
		return errors.New("The configuration contains a invalid frequency")
	}

	if config.Output == nil {
		return errors.New("The configuration did not have an output channel")
	}
	return nil
}

// DefaultConfig returns a config with values that have been tested
func DefaultConfig() *AmassConfig {
	config := &AmassConfig{
		Ports:           []int{443},
		Recursive:       true,
		Alterations:     true,
		Frequency:       50 * time.Millisecond,
		MinForRecursive: 1,
	}
	return config
}

// Ensures that all configuration elements have valid values
func CustomConfig(ac *AmassConfig) *AmassConfig {
	config := DefaultConfig()

	if len(ac.Domains()) > 0 {
		config.AddDomains(ac.Domains())
	}
	if len(ac.Ports) > 0 {
		config.Ports = ac.Ports
	}
	if len(ac.Wordlist) == 0 {
		config.Wordlist = GetDefaultWordlist()
	} else {
		config.Wordlist = ac.Wordlist
	}
	// Check that the config values have been set appropriately
	if ac.Frequency > config.Frequency {
		config.Frequency = ac.Frequency
	}
	if ac.proxy != nil {
		config.proxy = ac.proxy
	}
	if ac.MinForRecursive > config.MinForRecursive {
		config.MinForRecursive = ac.MinForRecursive
	}
	config.ASNs = ac.ASNs
	config.CIDRs = ac.CIDRs
	config.IPs = ac.IPs
	config.BruteForcing = ac.BruteForcing
	config.Recursive = ac.Recursive
	config.Alterations = ac.Alterations
	config.Output = ac.Output
	config.AdditionalDomains = ac.AdditionalDomains
	config.Resolvers = ac.Resolvers
	config.Blacklist = ac.Blacklist
	config.Setup()
	return config
}

func GetDefaultWordlist() []string {
	var list []string
	var wordlist io.Reader

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

	scanner := bufio.NewScanner(wordlist)
	// Once we have used all the words, we are finished
	for scanner.Scan() {
		// Get the next word in the list
		word := scanner.Text()
		if word != "" {
			// Add the word to the list
			list = append(list, word)
		}
	}
	return list
}

//--------------------------------------------------------------------------------------------------
// ReverseWhois - Returns domain names that are related to the domain provided
func (c *AmassConfig) ReverseWhois(domain string) []string {
	var domains []string

	page := GetWebPageWithDialContext(c.DialContext,
		"http://viewdns.info/reversewhois/?q="+domain)
	if page == "" {
		return []string{}
	}
	// Pull the table we need from the page content
	table := getViewDNSTable(page)
	// Get the list of domain names discovered through
	// the reverse DNS service
	re := regexp.MustCompile("
([a-zA-Z0-9]{1}[a-zA-Z0-9-]{0,61}[a-zA-Z0-9]{1}[.]{1}[a-zA-Z0-9-]+)	")
	subs := re.FindAllStringSubmatch(table, -1)
	for _, match := range subs {
		sub := match[1]
		if sub == "" {
			continue
		}
		domains = append(domains, strings.TrimSpace(sub))
	}
	sort.Strings(domains)
	return domains
}

func getViewDNSTable(page string) string {
	var begin, end int
	s := page

	for i := 0; i < 4; i++ {
		b := strings.Index(s, "<table")
		if b == -1 {
			return ""
		}
		begin += b + 6

		if e := strings.Index(s[b:], "
"); e == -1 {
			return ""
		} else {
			end = begin + e
		}

		s = page[end+8:]
	}

	i := strings.Index(page[begin:end], "<table")
	i = strings.Index(page[begin+i+6:end], "<table")
	return page[begin+i : end]
}

//--------------------------------------------------------------------------------------------------
// Methods that handle networking that is specific to the Amass configuration

func (c *AmassConfig) SetupProxyConnection(addr string) error {
	client, err := proxyclient.NewProxyClient(addr)
	if err == nil {
		c.proxy = client
		// Override the Go default DNS resolver to prevent leakage
		net.DefaultResolver = &net.Resolver{
			PreferGo: true,
			Dial:     c.DNSDialContext,
		}
	}
	return err
}

func (c *AmassConfig) DNSDialContext(ctx context.Context, network, address string) (net.Conn, error) {
	resolver := c.resolver.Next()

	if c.proxy != nil {
		if timeout, ok := ctx.Deadline(); ok {
			return c.proxy.DialTimeout(network, resolver, timeout.Sub(time.Now()))
		}
		return c.proxy.Dial(network, resolver)
	}

	d := &net.Dialer{}
	return d.DialContext(ctx, network, resolver)
}

func (c *AmassConfig) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	if c.proxy != nil {
		if timeout, ok := ctx.Deadline(); ok {
			return c.proxy.DialTimeout(network, address, timeout.Sub(time.Now()))
		}
		return c.proxy.Dial(network, address)
	}

	d := &net.Dialer{
		Resolver: &net.Resolver{
			PreferGo: true,
			Dial:     c.DNSDialContext,
		},
	}
	return d.DialContext(ctx, network, address)
}
<commit_msg>slightly increased the default DNS query frequency<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package amass

import (
	"bufio"
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/gamexg/proxyclient"
)

// AmassConfig - Passes along optional configurations
type AmassConfig struct {
	sync.Mutex

	// The channel that will receive the results
	Output chan *AmassRequest

	// The ASNs that the enumeration will target
	ASNs []int

	// The CIDRs that the enumeration will target
	CIDRs []*net.IPNet

	// The IPs that the enumeration will target
	IPs []net.IP

	// The ports that will be checked for certificates
	Ports []int

	// The list of words to use when generating names
	Wordlist []string

	// Will the enumeration including brute forcing techniques
	BruteForcing bool

	// Will recursive brute forcing be performed?
	Recursive bool

	// Minimum number of subdomain discoveries before performing recursive brute forcing
	MinForRecursive int

	// Will discovered subdomain name alterations be generated?
	Alterations bool

	// A blacklist of subdomain names that will not be investigated
	Blacklist []string

	// Sets the maximum number of DNS queries per minute
	Frequency time.Duration

	// Preferred DNS resolvers identified by the user
	Resolvers []string

	// Indicate that Amass cannot add domains to the config
	AdditionalDomains bool

	// The root domain names that the enumeration will target
	domains []string

	// Is responsible for performing simple DNS resolutions
	dns *queries

	// Handles selecting the next DNS resolver to be used
	resolver *resolvers

	// Performs lookups of root domain names from subdomain names
	domainLookup *DomainLookup

	// Detects DNS wildcards
	wildcards *Wildcards

	// The optional proxy connection for the enumeration to use
	proxy proxyclient.ProxyClient
}

func (c *AmassConfig) Setup() {
	// Setup the services potentially needed by all of amass
	c.dns = newQueriesSubsystem(c)
	c.domainLookup = NewDomainLookup(c)
	c.wildcards = NewWildcardDetection(c)
	c.resolver = newResolversSubsystem(c)
}

func (c *AmassConfig) AddDomains(names []string) {
	c.Lock()
	defer c.Unlock()

	c.domains = UniqueAppend(c.domains, names...)
}

func (c *AmassConfig) Domains() []string {
	c.Lock()
	defer c.Unlock()

	return c.domains
}

func (c *AmassConfig) Blacklisted(name string) bool {
	var resp bool

	for _, bl := range c.Blacklist {
		if match := strings.HasSuffix(name, bl); match {
			resp = true
			break
		}
	}
	return resp
}

func CheckConfig(config *AmassConfig) error {
	if len(config.Wordlist) == 0 {
		return errors.New("The configuration contains no wordlist")
	}

	if config.Frequency < DefaultConfig().Frequency {
		return errors.New("The configuration contains a invalid frequency")
	}

	if config.Output == nil {
		return errors.New("The configuration did not have an output channel")
	}
	return nil
}

// DefaultConfig returns a config with values that have been tested
func DefaultConfig() *AmassConfig {
	config := &AmassConfig{
		Ports:           []int{443},
		Recursive:       true,
		Alterations:     true,
		Frequency:       30 * time.Millisecond,
		MinForRecursive: 1,
	}
	return config
}

// Ensures that all configuration elements have valid values
func CustomConfig(ac *AmassConfig) *AmassConfig {
	config := DefaultConfig()

	if len(ac.Domains()) > 0 {
		config.AddDomains(ac.Domains())
	}
	if len(ac.Ports) > 0 {
		config.Ports = ac.Ports
	}
	if len(ac.Wordlist) == 0 {
		config.Wordlist = GetDefaultWordlist()
	} else {
		config.Wordlist = ac.Wordlist
	}
	// Check that the config values have been set appropriately
	if ac.Frequency > config.Frequency {
		config.Frequency = ac.Frequency
	}
	if ac.proxy != nil {
		config.proxy = ac.proxy
	}
	if ac.MinForRecursive > config.MinForRecursive {
		config.MinForRecursive = ac.MinForRecursive
	}
	config.ASNs = ac.ASNs
	config.CIDRs = ac.CIDRs
	config.IPs = ac.IPs
	config.BruteForcing = ac.BruteForcing
	config.Recursive = ac.Recursive
	config.Alterations = ac.Alterations
	config.Output = ac.Output
	config.AdditionalDomains = ac.AdditionalDomains
	config.Resolvers = ac.Resolvers
	config.Blacklist = ac.Blacklist
	config.Setup()
	return config
}

func GetDefaultWordlist() []string {
	var list []string
	var wordlist io.Reader

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

	scanner := bufio.NewScanner(wordlist)
	// Once we have used all the words, we are finished
	for scanner.Scan() {
		// Get the next word in the list
		word := scanner.Text()
		if word != "" {
			// Add the word to the list
			list = append(list, word)
		}
	}
	return list
}

//--------------------------------------------------------------------------------------------------
// ReverseWhois - Returns domain names that are related to the domain provided
func (c *AmassConfig) ReverseWhois(domain string) []string {
	var domains []string

	page := GetWebPageWithDialContext(c.DialContext,
		"http://viewdns.info/reversewhois/?q="+domain)
	if page == "" {
		return []string{}
	}
	// Pull the table we need from the page content
	table := getViewDNSTable(page)
	// Get the list of domain names discovered through
	// the reverse DNS service
	re := regexp.MustCompile("
([a-zA-Z0-9]{1}[a-zA-Z0-9-]{0,61}[a-zA-Z0-9]{1}[.]{1}[a-zA-Z0-9-]+)	")
	subs := re.FindAllStringSubmatch(table, -1)
	for _, match := range subs {
		sub := match[1]
		if sub == "" {
			continue
		}
		domains = append(domains, strings.TrimSpace(sub))
	}
	sort.Strings(domains)
	return domains
}

func getViewDNSTable(page string) string {
	var begin, end int
	s := page

	for i := 0; i < 4; i++ {
		b := strings.Index(s, "<table")
		if b == -1 {
			return ""
		}
		begin += b + 6

		if e := strings.Index(s[b:], "
"); e == -1 {
			return ""
		} else {
			end = begin + e
		}

		s = page[end+8:]
	}

	i := strings.Index(page[begin:end], "<table")
	i = strings.Index(page[begin+i+6:end], "<table")
	return page[begin+i : end]
}

//--------------------------------------------------------------------------------------------------
// Methods that handle networking that is specific to the Amass configuration

func (c *AmassConfig) SetupProxyConnection(addr string) error {
	client, err := proxyclient.NewProxyClient(addr)
	if err == nil {
		c.proxy = client
		// Override the Go default DNS resolver to prevent leakage
		net.DefaultResolver = &net.Resolver{
			PreferGo: true,
			Dial:     c.DNSDialContext,
		}
	}
	return err
}

func (c *AmassConfig) DNSDialContext(ctx context.Context, network, address string) (net.Conn, error) {
	resolver := c.resolver.Next()

	if c.proxy != nil {
		if timeout, ok := ctx.Deadline(); ok {
			return c.proxy.DialTimeout(network, resolver, timeout.Sub(time.Now()))
		}
		return c.proxy.Dial(network, resolver)
	}

	d := &net.Dialer{}
	return d.DialContext(ctx, network, resolver)
}

func (c *AmassConfig) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	if c.proxy != nil {
		if timeout, ok := ctx.Deadline(); ok {
			return c.proxy.DialTimeout(network, address, timeout.Sub(time.Now()))
		}
		return c.proxy.Dial(network, address)
	}

	d := &net.Dialer{
		Resolver: &net.Resolver{
			PreferGo: true,
			Dial:     c.DNSDialContext,
		},
	}
	return d.DialContext(ctx, network, address)
}
<|endoftext|>


<commit_before>//
// Copyright © 2011-2017 Guy M. Allard
//
// 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 stompngo

import (
	"bufio"
	"bytes"
	"net"
	// "bytes"
	"strconv"
	"time"
)

/*
	Logical network writer.  Read wiredata structures from the communication
	channel, and put the frame on the wire.
*/
func (c *Connection) writer() {
writerLoop:
	for {
		select {
		case d := <-c.output:
			c.log("WTR_WIREWRITE start")
			c.wireWrite(d)
			c.log("WTR_WIREWRITE COMPLETE", d.frame.Command, d.frame.Headers,
				HexData(d.frame.Body))
			if d.frame.Command == DISCONNECT {
				break writerLoop // we are done with this connection
			}
		case _ = <-c.ssdc:
			c.log("WTR_WIREWRITE shutdown S received")
			break writerLoop
		case _ = <-c.wtrsdc:
			c.log("WTR_WIREWRITE shutdown W received")
			break writerLoop
		}
	} // of for
	//
	c.connected = false
	c.log("WTR_SHUTDOWN", time.Now())
}

/*
	Connection logical write.
*/
func (c *Connection) wireWrite(d wiredata) {
	f := &d.frame
	// fmt.Printf("WWD01 f:[%v]\n", f)
	switch f.Command {
	case "\n": // HeartBeat frame
		if c.dld.wde && c.dld.dns {
			_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
		}
		_, e := c.wtr.WriteString(f.Command)
		if e != nil {
			if e.(net.Error).Timeout() {
				if c.dld.dns {
					c.dld.dlnotify(e, true)
				}
			}
			d.errchan <- e
			return
		}
	default: // Other frames
		if e := f.writeFrame(c.wtr, c); e != nil {
			d.errchan <- e
			return
		}
		if e := c.wtr.Flush(); e != nil {
			d.errchan <- e
			return
		}
	}
	if e := c.wtr.Flush(); e != nil {
		d.errchan <- e
		return
	}
	//
	if c.hbd != nil {
		c.hbd.sdl.Lock()
		c.hbd.ls = time.Now().UnixNano() // Latest good send
		c.hbd.sdl.Unlock()
	}
	c.mets.tfw++                // Frame written count
	c.mets.tbw += f.Size(false) // Bytes written count
	//
	d.errchan <- nil
	return
}

/*
	Physical frame write to the wire.
*/
func (f *Frame) writeFrame(w *bufio.Writer, c *Connection) error {

	var sctok bool
	// Content type.  Always add it if the client does not suppress and does not
	// supply it.
	_, sctok = f.Headers.Contains(HK_SUPPRESS_CT)
	if !sctok {
		if _, ctok := f.Headers.Contains(HK_CONTENT_TYPE); !ctok {
			f.Headers = append(f.Headers, HK_CONTENT_TYPE,
				DFLT_CONTENT_TYPE)
		}
	}

	var sclok bool
	// Content length - Always add it if client does not suppress it and
	// does not supply it.
	_, sclok = f.Headers.Contains(HK_SUPPRESS_CL)
	if !sclok {
		if _, clok := f.Headers.Contains(HK_CONTENT_LENGTH); !clok {
			f.Headers = append(f.Headers, HK_CONTENT_LENGTH, strconv.Itoa(len(f.Body)))
		}
	}
	// Encode the headers if needed
	if c.Protocol() > SPL_10 && f.Command != CONNECT {
		for i := 0; i < len(f.Headers); i += 2 {
			f.Headers[i] = encode(f.Headers[i])
			f.Headers[i+1] = encode(f.Headers[i+1])
		}
	}

	if sclok {
		nz := bytes.IndexByte(f.Body, 0)
		// fmt.Printf("WDBG41 ok:%v\n", nz)
		if nz == 0 {
			f.Body = []byte{}
			// fmt.Printf("WDBG42 body:%v bodystring: %v\n", f.Body, string(f.Body))
		} else if nz > 0 {
			f.Body = f.Body[0:nz]
			// fmt.Printf("WDBG43 body:%v bodystring: %v\n", f.Body, string(f.Body))
		}
	}

	if c.dld.wde && c.dld.dns {
		_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
	}

	// Writes start

	// Write the frame Command
	_, e := w.WriteString(f.Command + "\n")
	if c.checkWriteError(e) != nil {
		return e
	}
	// fmt.Println("WRCMD", f.Command)
	// Write the frame Headers
	for i := 0; i < len(f.Headers); i += 2 {
		if c.dld.wde && c.dld.dns {
			_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
		}
		_, e := w.WriteString(f.Headers[i] + ":" + f.Headers[i+1] + "\n")
		if c.checkWriteError(e) != nil {
			return e
		}
		// fmt.Println("WRHDR", f.Headers[i]+":"+f.Headers[i+1]+"\n")
	}

	// Write the last Header LF
	if c.dld.wde && c.dld.dns {
		_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
	}
	e = w.WriteByte('\n')
	if c.checkWriteError(e) != nil {
		return e
	}
	// fmt.Printf("WDBG40 ok:%v\n", sclok)

	// Write the body
	if len(f.Body) != 0 { // Foolish to write 0 length data
		// fmt.Println("WRBDY", f.Body)
		e := c.writeBody(f)
		if c.checkWriteError(e) != nil {
			return e
		}
	}
	if c.dld.wde && c.dld.dns {
		_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
	}
	e = w.WriteByte(0)
	if c.checkWriteError(e) != nil {
		return e
	}
	// End of write loop - set no deadline
	if c.dld.wde {
		_ = c.netconn.SetWriteDeadline(c.dld.t0)
	}
	return nil
}

func (c *Connection) checkWriteError(e error) error {
	if e == nil {
		return e
	}
	ne, ok := e.(net.Error)
	if !ok {
		return e
	}
	if ne.Timeout() {
		if c.dld.dns {
			c.log("invoking write deadline callback 1")
			c.dld.dlnotify(e, true)
		}
	}
	return e
}

func (c *Connection) writeBody(f *Frame) error {
	// fmt.Printf("WDBG99 body:%v bodystring: %v\n", f.Body, string(f.Body))
	var n = 0
	var e error
	for {
		if c.dld.wde && c.dld.dns {
			_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
		}
		n, e = c.wtr.Write(f.Body)
		if n == len(f.Body) {
			return e
		}
		c.log("SHORT WRITE", n, len(f.Body))
		if n == 0 { // Zero bytes would mean something is seriously wrong.
			return e
		}
		if !c.dld.rfsw {
			return e
		}
		if c.dld.wde && c.dld.dns && isErrorTimeout(e) {
			c.log("invoking write deadline callback 2")
			c.dld.dlnotify(e, true)
		}
		// *Any* error from a bufio.Writer is *not* recoverable.  See code in
		// bufio.go to understand this.  We get a new writer here, to clear any
		// error condition.
		c.wtr = bufio.NewWriter(c.netconn) // Create new writer
		f.Body = f.Body[n:]
	}
}

func isErrorTimeout(e error) bool {
	if e == nil {
		return false
	}
	_, ok := e.(net.Error)
	if !ok {
		return false
	}
	return true
}
<commit_msg>Issue #38  fix WriteDeadline to make it work as expected<commit_after>//
// Copyright © 2011-2017 Guy M. Allard
//
// 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 stompngo

import (
	"bufio"
	"bytes"
	"net"
	// "bytes"
	"strconv"
	"time"
)

/*
	Logical network writer.  Read wiredata structures from the communication
	channel, and put the frame on the wire.
*/
func (c *Connection) writer() {
writerLoop:
	for {
		select {
		case d := <-c.output:
			c.log("WTR_WIREWRITE start")
			c.wireWrite(d)
			c.log("WTR_WIREWRITE COMPLETE", d.frame.Command, d.frame.Headers,
				HexData(d.frame.Body))
			if d.frame.Command == DISCONNECT {
				break writerLoop // we are done with this connection
			}
		case _ = <-c.ssdc:
			c.log("WTR_WIREWRITE shutdown S received")
			break writerLoop
		case _ = <-c.wtrsdc:
			c.log("WTR_WIREWRITE shutdown W received")
			break writerLoop
		}
	} // of for
	//
	c.connected = false
	c.log("WTR_SHUTDOWN", time.Now())
}

/*
	Connection logical write.
*/
func (c *Connection) wireWrite(d wiredata) {
	f := &d.frame
	// fmt.Printf("WWD01 f:[%v]\n", f)
	switch f.Command {
	case "\n": // HeartBeat frame
		if c.dld.wde && c.dld.wds {
			_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
		}
		_, e := c.wtr.WriteString(f.Command)
		if e != nil {
			if e.(net.Error).Timeout() {
				if c.dld.dns {
					c.dld.dlnotify(e, true)
				}
			}
			d.errchan <- e
			return
		}
	default: // Other frames
		if e := f.writeFrame(c.wtr, c); e != nil {
			d.errchan <- e
			return
		}
		if e := c.wtr.Flush(); e != nil {
			d.errchan <- e
			return
		}
	}
	if e := c.wtr.Flush(); e != nil {
		d.errchan <- e
		return
	}
	//
	if c.hbd != nil {
		c.hbd.sdl.Lock()
		c.hbd.ls = time.Now().UnixNano() // Latest good send
		c.hbd.sdl.Unlock()
	}
	c.mets.tfw++                // Frame written count
	c.mets.tbw += f.Size(false) // Bytes written count
	//
	d.errchan <- nil
	return
}

/*
	Physical frame write to the wire.
*/
func (f *Frame) writeFrame(w *bufio.Writer, c *Connection) error {

	var sctok bool
	// Content type.  Always add it if the client does not suppress and does not
	// supply it.
	_, sctok = f.Headers.Contains(HK_SUPPRESS_CT)
	if !sctok {
		if _, ctok := f.Headers.Contains(HK_CONTENT_TYPE); !ctok {
			f.Headers = append(f.Headers, HK_CONTENT_TYPE,
				DFLT_CONTENT_TYPE)
		}
	}

	var sclok bool
	// Content length - Always add it if client does not suppress it and
	// does not supply it.
	_, sclok = f.Headers.Contains(HK_SUPPRESS_CL)
	if !sclok {
		if _, clok := f.Headers.Contains(HK_CONTENT_LENGTH); !clok {
			f.Headers = append(f.Headers, HK_CONTENT_LENGTH, strconv.Itoa(len(f.Body)))
		}
	}
	// Encode the headers if needed
	if c.Protocol() > SPL_10 && f.Command != CONNECT {
		for i := 0; i < len(f.Headers); i += 2 {
			f.Headers[i] = encode(f.Headers[i])
			f.Headers[i+1] = encode(f.Headers[i+1])
		}
	}

	if sclok {
		nz := bytes.IndexByte(f.Body, 0)
		// fmt.Printf("WDBG41 ok:%v\n", nz)
		if nz == 0 {
			f.Body = []byte{}
			// fmt.Printf("WDBG42 body:%v bodystring: %v\n", f.Body, string(f.Body))
		} else if nz > 0 {
			f.Body = f.Body[0:nz]
			// fmt.Printf("WDBG43 body:%v bodystring: %v\n", f.Body, string(f.Body))
		}
	}

	if c.dld.wde && c.dld.wds {
		_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
	}

	// Writes start

	// Write the frame Command
	_, e := w.WriteString(f.Command + "\n")
	if c.checkWriteError(e) != nil {
		return e
	}
	// fmt.Println("WRCMD", f.Command)
	// Write the frame Headers
	for i := 0; i < len(f.Headers); i += 2 {
		if c.dld.wde && c.dld.wds {
			_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
		}
		_, e := w.WriteString(f.Headers[i] + ":" + f.Headers[i+1] + "\n")
		if c.checkWriteError(e) != nil {
			return e
		}
		// fmt.Println("WRHDR", f.Headers[i]+":"+f.Headers[i+1]+"\n")
	}

	// Write the last Header LF
	if c.dld.wde && c.dld.wds {
		_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
	}
	e = w.WriteByte('\n')
	if c.checkWriteError(e) != nil {
		return e
	}
	// fmt.Printf("WDBG40 ok:%v\n", sclok)

	// Write the body
	if len(f.Body) != 0 { // Foolish to write 0 length data
		// fmt.Println("WRBDY", f.Body)
		e := c.writeBody(f)
		if c.checkWriteError(e) != nil {
			return e
		}
	}
	if c.dld.wde && c.dld.wds {
		_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
	}
	e = w.WriteByte(0)
	if c.checkWriteError(e) != nil {
		return e
	}
	// End of write loop - set no deadline
	if c.dld.wde {
		_ = c.netconn.SetWriteDeadline(c.dld.t0)
	}
	return nil
}

func (c *Connection) checkWriteError(e error) error {
	if e == nil {
		return e
	}
	ne, ok := e.(net.Error)
	if !ok {
		return e
	}
	if ne.Timeout() {
		if c.dld.dns {
			c.log("invoking write deadline callback 1")
			c.dld.dlnotify(e, true)
		}
	}
	return e
}

func (c *Connection) writeBody(f *Frame) error {
	// fmt.Printf("WDBG99 body:%v bodystring: %v\n", f.Body, string(f.Body))
	var n = 0
	var e error
	for {
		if c.dld.wde && c.dld.wds {
			_ = c.netconn.SetWriteDeadline(time.Now().Add(c.dld.wdld))
		}
		n, e = c.wtr.Write(f.Body)
		if n == len(f.Body) {
			return e
		}
		c.log("SHORT WRITE", n, len(f.Body))
		if n == 0 { // Zero bytes would mean something is seriously wrong.
			return e
		}
		if !c.dld.rfsw {
			return e
		}
		if c.dld.wde && c.dld.wds && isErrorTimeout(e) {
			c.log("invoking write deadline callback 2")
			c.dld.dlnotify(e, true)
		}
		// *Any* error from a bufio.Writer is *not* recoverable.  See code in
		// bufio.go to understand this.  We get a new writer here, to clear any
		// error condition.
		c.wtr = bufio.NewWriter(c.netconn) // Create new writer
		f.Body = f.Body[n:]
	}
}

func isErrorTimeout(e error) bool {
	if e == nil {
		return false
	}
	_, ok := e.(net.Error)
	if !ok {
		return false
	}
	return true
}
<|endoftext|>