<commit_before>// Copyright 2013 Prometheus Team
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/prometheus/client_golang/extraction"

	clientmodel "github.com/prometheus/client_golang/model"

	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/notification"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/rules"
	"github.com/prometheus/prometheus/storage/metric"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/prometheus/prometheus/storage/remote/opentsdb"
	"github.com/prometheus/prometheus/web"
	"github.com/prometheus/prometheus/web/api"
)

const deletionBatchSize = 100

// Commandline flags.
var (
	configFile         = flag.String("configFile", "prometheus.conf", "Prometheus configuration file name.")
	metricsStoragePath = flag.String("metricsStoragePath", "/tmp/metrics", "Base path for metrics storage.")

	alertmanagerUrl = flag.String("alertmanager.url", "", "The URL of the alert manager to send notifications to.")

	remoteTSDBUrl     = flag.String("storage.remote.url", "", "The URL of the OpenTSDB instance to send samples to.")
	remoteTSDBTimeout = flag.Duration("storage.remote.timeout", 30*time.Second, "The timeout to use when sending samples to OpenTSDB.")

	samplesQueueCapacity      = flag.Int("storage.queue.samplesCapacity", 4096, "The size of the unwritten samples queue.")
	diskAppendQueueCapacity   = flag.Int("storage.queue.diskAppendCapacity", 1000000, "The size of the queue for items that are pending writing to disk.")
	memoryAppendQueueCapacity = flag.Int("storage.queue.memoryAppendCapacity", 10000, "The size of the queue for items that are pending writing to memory.")

	compactInterval = flag.Duration("compact.interval", 3*time.Hour, "The amount of time between compactions.")
	compactGroupSize = flag.Int("compact.groupSize", 500, "The minimum group size for compacted samples.")
	compactAgeInclusiveness = flag.Duration("compact.ageInclusiveness", 5*time.Minute, "The age beyond which samples should be compacted.")

	deleteInterval = flag.Duration("delete.interval", 11*time.Hour, "The amount of time between deletion of old values.")

	deleteAge = flag.Duration("delete.ageMaximum", 15*24*time.Hour, "The relative maximum age for values before they are deleted.")

	arenaFlushInterval = flag.Duration("arena.flushInterval", 15*time.Minute, "The period at which the in-memory arena is flushed to disk.")
	arenaTTL           = flag.Duration("arena.ttl", 10*time.Minute, "The relative age of values to purge to disk from memory.")

	notificationQueueCapacity = flag.Int("alertmanager.notificationQueueCapacity", 100, "The size of the queue for pending alert manager notifications.")

	concurrentRetrievalAllowance = flag.Int("concurrentRetrievalAllowance", 15, "The number of concurrent metrics retrieval requests allowed.")

	printVersion = flag.Bool("version", false, "print version information")
)

type prometheus struct {
	compactionTimer *time.Ticker
	deletionTimer       *time.Ticker

	curationSema             chan bool
	stopBackgroundOperations chan bool

	unwrittenSamples chan *extraction.Result

	ruleManager     rules.RuleManager
	targetManager   retrieval.TargetManager
	notifications   chan notification.NotificationReqs
	storage         *metric.TieredStorage
	remoteTSDBQueue *remote.TSDBQueueManager

	curationState metric.CurationStateUpdater
}

func (p *prometheus) interruptHandler() {
	notifier := make(chan os.Signal)
	signal.Notify(notifier, os.Interrupt, syscall.SIGTERM)

	<-notifier

	glog.Warning("Received SIGINT/SIGTERM; Exiting gracefully...")
	p.close()
	os.Exit(0)
}

func (p *prometheus) compact(olderThan time.Duration, groupSize int) error {
	select {
	case p.curationSema <- true:
	default:
		glog.Warningf("Deferred compaction for %s and %s due to existing operation.", olderThan, groupSize)

		return nil
	}

	defer func() {
		<-p.curationSema
	}()

	processor := metric.NewCompactionProcessor(&metric.CompactionProcessorOptions{
		MaximumMutationPoolBatch: groupSize * 3,
		MinimumGroupSize:         groupSize,
	})
	defer processor.Close()

	curator := metric.NewCurator(&metric.CuratorOptions{
		Stop: p.stopBackgroundOperations,

		ViewQueue: p.storage.ViewQueue,
	})
	defer curator.Close()

	return curator.Run(olderThan, clientmodel.Now(), processor, p.storage.DiskStorage.CurationRemarks, p.storage.DiskStorage.MetricSamples, p.storage.DiskStorage.MetricHighWatermarks, p.curationState)
}

func (p *prometheus) delete(olderThan time.Duration, batchSize int) error {
	select {
	case p.curationSema <- true:
	default:
		glog.Warningf("Deferred deletion for %s due to existing operation.", olderThan)

		return nil
	}

	processor := metric.NewDeletionProcessor(&metric.DeletionProcessorOptions{
		MaximumMutationPoolBatch: batchSize,
	})
	defer processor.Close()

	curator := metric.NewCurator(&metric.CuratorOptions{
		Stop: p.stopBackgroundOperations,

		ViewQueue: p.storage.ViewQueue,
	})
	defer curator.Close()

	return curator.Run(olderThan, clientmodel.Now(), processor, p.storage.DiskStorage.CurationRemarks, p.storage.DiskStorage.MetricSamples, p.storage.DiskStorage.MetricHighWatermarks, p.curationState)
}

func (p *prometheus) close() {
	select {
	case p.curationSema <- true:
	default:
	}

	if p.compactionTimer != nil {
		p.compactionTimer.Stop()
	}
	if p.deletionTimer != nil {
		p.deletionTimer.Stop()
	}

	// Stop any currently active curation (deletion or compaction).
	if len(p.stopBackgroundOperations) == 0 {
		p.stopBackgroundOperations <- true
	}

	p.ruleManager.Stop()
	p.targetManager.Stop()
	close(p.unwrittenSamples)

	p.storage.Close()

	if p.remoteTSDBQueue != nil {
		p.remoteTSDBQueue.Close()
	}

	close(p.notifications)
	close(p.stopBackgroundOperations)
}

func main() {
	// TODO(all): Future additions to main should be, where applicable, glumped
	// into the prometheus struct above---at least where the scoping of the entire
	// server is concerned.
	flag.Parse()

	versionInfoTmpl.Execute(os.Stdout, BuildInfo)

	if *printVersion {
		os.Exit(0)
	}

	conf, err := config.LoadFromFile(*configFile)
	if err != nil {
		glog.Fatalf("Error loading configuration from %s: %v", *configFile, err)
	}

	ts, err := metric.NewTieredStorage(uint(*diskAppendQueueCapacity), 100, *arenaFlushInterval, *arenaTTL, *metricsStoragePath)
	if err != nil {
		glog.Fatal("Error opening storage: ", err)
	}

	var remoteTSDBQueue *remote.TSDBQueueManager = nil
	if *remoteTSDBUrl == "" {
		glog.Warningf("No TSDB URL provided; not sending any samples to long-term storage")
	} else {
		openTSDB := opentsdb.NewClient(*remoteTSDBUrl, *remoteTSDBTimeout)
		remoteTSDBQueue = remote.NewTSDBQueueManager(openTSDB, 512)
		go remoteTSDBQueue.Run()
	}

	unwrittenSamples := make(chan *extraction.Result, *samplesQueueCapacity)
	ingester := &retrieval.MergeLabelsIngester{
		Labels:          conf.GlobalLabels(),
		CollisionPrefix: clientmodel.ExporterLabelPrefix,

		Ingester: retrieval.ChannelIngester(unwrittenSamples),
	}

	compactionTimer := time.NewTicker(*compactInterval)
	deletionTimer := time.NewTicker(*deleteInterval)

	// Queue depth will need to be exposed
	targetManager := retrieval.NewTargetManager(ingester, *concurrentRetrievalAllowance)
	targetManager.AddTargetsFromConfig(conf)

	notifications := make(chan notification.NotificationReqs, *notificationQueueCapacity)

	// Queue depth will need to be exposed
	ruleManager := rules.NewRuleManager(&rules.RuleManagerOptions{
		Results:            unwrittenSamples,
		Notifications:      notifications,
		EvaluationInterval: conf.EvaluationInterval(),
		Storage:            ts,
		PrometheusUrl:      web.MustBuildServerUrl(),
	})
	if err := ruleManager.AddRulesFromConfig(conf); err != nil {
		glog.Fatal("Error loading rule files: ", err)
	}
	go ruleManager.Run()

	notificationHandler := notification.NewNotificationHandler(*alertmanagerUrl, notifications)
	go notificationHandler.Run()

	flags := map[string]string{}

	flag.VisitAll(func(f *flag.Flag) {
		flags[f.Name] = f.Value.String()
	})

	prometheusStatus := &web.PrometheusStatusHandler{
		BuildInfo:   BuildInfo,
		Config:      conf.String(),
		RuleManager: ruleManager,
		TargetPools: targetManager.Pools(),
		Flags:       flags,
		Birth:       time.Now(),
	}

	alertsHandler := &web.AlertsHandler{
		RuleManager: ruleManager,
	}

	databasesHandler := &web.DatabasesHandler{
		Provider:        ts.DiskStorage,
		RefreshInterval: 5 * time.Minute,
	}

	metricsService := &api.MetricsService{
		Config:        &conf,
		TargetManager: targetManager,
		Storage:       ts,
	}

	webService := &web.WebService{
		StatusHandler:    prometheusStatus,
		MetricsHandler:   metricsService,
		DatabasesHandler: databasesHandler,
		AlertsHandler:    alertsHandler,
	}

	prometheus := &prometheus{
		compactionTimer: compactionTimer,

		deletionTimer: deletionTimer,

		curationState: prometheusStatus,
		curationSema:  make(chan bool, 1),

		unwrittenSamples: unwrittenSamples,

		stopBackgroundOperations: make(chan bool, 1),

		ruleManager:     ruleManager,
		targetManager:   targetManager,
		notifications:   notifications,
		storage:         ts,
		remoteTSDBQueue: remoteTSDBQueue,
	}
	defer prometheus.close()

	storageStarted := make(chan bool)
	go ts.Serve(storageStarted)
	<-storageStarted

	go prometheus.interruptHandler()

	go func() {
		for _ = range prometheus.compactionTimer.C {
			glog.Info("Starting compaction...")
			err := prometheus.compact(*compactAgeInclusiveness, *compactGroupSize)

			if err != nil {
				glog.Error("could not compact: ", err)
			}
			glog.Info("Done")
		}
	}()

	go func() {
		for _ = range prometheus.deletionTimer.C {
			glog.Info("Starting deletion of stale values...")
			err := prometheus.delete(*deleteAge, deletionBatchSize)

			if err != nil {
				glog.Error("could not delete: ", err)
			}
			glog.Info("Done")
		}
	}()

	go func() {
		err := webService.ServeForever()
		if err != nil {
			glog.Fatal(err)
		}
	}()

	// TODO(all): Migrate this into prometheus.serve().
	for block := range unwrittenSamples {
		if block.Err == nil && len(block.Samples) > 0 {
			ts.AppendSamples(block.Samples)
			if remoteTSDBQueue != nil {
				remoteTSDBQueue.Queue(block.Samples)
			}
		}
	}
}
<commit_msg>"go fmt" fixup.<commit_after>// Copyright 2013 Prometheus Team
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/golang/glog"
	"github.com/prometheus/client_golang/extraction"

	clientmodel "github.com/prometheus/client_golang/model"

	"github.com/prometheus/prometheus/config"
	"github.com/prometheus/prometheus/notification"
	"github.com/prometheus/prometheus/retrieval"
	"github.com/prometheus/prometheus/rules"
	"github.com/prometheus/prometheus/storage/metric"
	"github.com/prometheus/prometheus/storage/remote"
	"github.com/prometheus/prometheus/storage/remote/opentsdb"
	"github.com/prometheus/prometheus/web"
	"github.com/prometheus/prometheus/web/api"
)

const deletionBatchSize = 100

// Commandline flags.
var (
	configFile         = flag.String("configFile", "prometheus.conf", "Prometheus configuration file name.")
	metricsStoragePath = flag.String("metricsStoragePath", "/tmp/metrics", "Base path for metrics storage.")

	alertmanagerUrl = flag.String("alertmanager.url", "", "The URL of the alert manager to send notifications to.")

	remoteTSDBUrl     = flag.String("storage.remote.url", "", "The URL of the OpenTSDB instance to send samples to.")
	remoteTSDBTimeout = flag.Duration("storage.remote.timeout", 30*time.Second, "The timeout to use when sending samples to OpenTSDB.")

	samplesQueueCapacity      = flag.Int("storage.queue.samplesCapacity", 4096, "The size of the unwritten samples queue.")
	diskAppendQueueCapacity   = flag.Int("storage.queue.diskAppendCapacity", 1000000, "The size of the queue for items that are pending writing to disk.")
	memoryAppendQueueCapacity = flag.Int("storage.queue.memoryAppendCapacity", 10000, "The size of the queue for items that are pending writing to memory.")

	compactInterval         = flag.Duration("compact.interval", 3*time.Hour, "The amount of time between compactions.")
	compactGroupSize        = flag.Int("compact.groupSize", 500, "The minimum group size for compacted samples.")
	compactAgeInclusiveness = flag.Duration("compact.ageInclusiveness", 5*time.Minute, "The age beyond which samples should be compacted.")

	deleteInterval = flag.Duration("delete.interval", 11*time.Hour, "The amount of time between deletion of old values.")

	deleteAge = flag.Duration("delete.ageMaximum", 15*24*time.Hour, "The relative maximum age for values before they are deleted.")

	arenaFlushInterval = flag.Duration("arena.flushInterval", 15*time.Minute, "The period at which the in-memory arena is flushed to disk.")
	arenaTTL           = flag.Duration("arena.ttl", 10*time.Minute, "The relative age of values to purge to disk from memory.")

	notificationQueueCapacity = flag.Int("alertmanager.notificationQueueCapacity", 100, "The size of the queue for pending alert manager notifications.")

	concurrentRetrievalAllowance = flag.Int("concurrentRetrievalAllowance", 15, "The number of concurrent metrics retrieval requests allowed.")

	printVersion = flag.Bool("version", false, "print version information")
)

type prometheus struct {
	compactionTimer *time.Ticker
	deletionTimer   *time.Ticker

	curationSema             chan bool
	stopBackgroundOperations chan bool

	unwrittenSamples chan *extraction.Result

	ruleManager     rules.RuleManager
	targetManager   retrieval.TargetManager
	notifications   chan notification.NotificationReqs
	storage         *metric.TieredStorage
	remoteTSDBQueue *remote.TSDBQueueManager

	curationState metric.CurationStateUpdater
}

func (p *prometheus) interruptHandler() {
	notifier := make(chan os.Signal)
	signal.Notify(notifier, os.Interrupt, syscall.SIGTERM)

	<-notifier

	glog.Warning("Received SIGINT/SIGTERM; Exiting gracefully...")
	p.close()
	os.Exit(0)
}

func (p *prometheus) compact(olderThan time.Duration, groupSize int) error {
	select {
	case p.curationSema <- true:
	default:
		glog.Warningf("Deferred compaction for %s and %s due to existing operation.", olderThan, groupSize)

		return nil
	}

	defer func() {
		<-p.curationSema
	}()

	processor := metric.NewCompactionProcessor(&metric.CompactionProcessorOptions{
		MaximumMutationPoolBatch: groupSize * 3,
		MinimumGroupSize:         groupSize,
	})
	defer processor.Close()

	curator := metric.NewCurator(&metric.CuratorOptions{
		Stop: p.stopBackgroundOperations,

		ViewQueue: p.storage.ViewQueue,
	})
	defer curator.Close()

	return curator.Run(olderThan, clientmodel.Now(), processor, p.storage.DiskStorage.CurationRemarks, p.storage.DiskStorage.MetricSamples, p.storage.DiskStorage.MetricHighWatermarks, p.curationState)
}

func (p *prometheus) delete(olderThan time.Duration, batchSize int) error {
	select {
	case p.curationSema <- true:
	default:
		glog.Warningf("Deferred deletion for %s due to existing operation.", olderThan)

		return nil
	}

	processor := metric.NewDeletionProcessor(&metric.DeletionProcessorOptions{
		MaximumMutationPoolBatch: batchSize,
	})
	defer processor.Close()

	curator := metric.NewCurator(&metric.CuratorOptions{
		Stop: p.stopBackgroundOperations,

		ViewQueue: p.storage.ViewQueue,
	})
	defer curator.Close()

	return curator.Run(olderThan, clientmodel.Now(), processor, p.storage.DiskStorage.CurationRemarks, p.storage.DiskStorage.MetricSamples, p.storage.DiskStorage.MetricHighWatermarks, p.curationState)
}

func (p *prometheus) close() {
	select {
	case p.curationSema <- true:
	default:
	}

	if p.compactionTimer != nil {
		p.compactionTimer.Stop()
	}
	if p.deletionTimer != nil {
		p.deletionTimer.Stop()
	}

	// Stop any currently active curation (deletion or compaction).
	if len(p.stopBackgroundOperations) == 0 {
		p.stopBackgroundOperations <- true
	}

	p.ruleManager.Stop()
	p.targetManager.Stop()
	close(p.unwrittenSamples)

	p.storage.Close()

	if p.remoteTSDBQueue != nil {
		p.remoteTSDBQueue.Close()
	}

	close(p.notifications)
	close(p.stopBackgroundOperations)
}

func main() {
	// TODO(all): Future additions to main should be, where applicable, glumped
	// into the prometheus struct above---at least where the scoping of the entire
	// server is concerned.
	flag.Parse()

	versionInfoTmpl.Execute(os.Stdout, BuildInfo)

	if *printVersion {
		os.Exit(0)
	}

	conf, err := config.LoadFromFile(*configFile)
	if err != nil {
		glog.Fatalf("Error loading configuration from %s: %v", *configFile, err)
	}

	ts, err := metric.NewTieredStorage(uint(*diskAppendQueueCapacity), 100, *arenaFlushInterval, *arenaTTL, *metricsStoragePath)
	if err != nil {
		glog.Fatal("Error opening storage: ", err)
	}

	var remoteTSDBQueue *remote.TSDBQueueManager = nil
	if *remoteTSDBUrl == "" {
		glog.Warningf("No TSDB URL provided; not sending any samples to long-term storage")
	} else {
		openTSDB := opentsdb.NewClient(*remoteTSDBUrl, *remoteTSDBTimeout)
		remoteTSDBQueue = remote.NewTSDBQueueManager(openTSDB, 512)
		go remoteTSDBQueue.Run()
	}

	unwrittenSamples := make(chan *extraction.Result, *samplesQueueCapacity)
	ingester := &retrieval.MergeLabelsIngester{
		Labels:          conf.GlobalLabels(),
		CollisionPrefix: clientmodel.ExporterLabelPrefix,

		Ingester: retrieval.ChannelIngester(unwrittenSamples),
	}

	compactionTimer := time.NewTicker(*compactInterval)
	deletionTimer := time.NewTicker(*deleteInterval)

	// Queue depth will need to be exposed
	targetManager := retrieval.NewTargetManager(ingester, *concurrentRetrievalAllowance)
	targetManager.AddTargetsFromConfig(conf)

	notifications := make(chan notification.NotificationReqs, *notificationQueueCapacity)

	// Queue depth will need to be exposed
	ruleManager := rules.NewRuleManager(&rules.RuleManagerOptions{
		Results:            unwrittenSamples,
		Notifications:      notifications,
		EvaluationInterval: conf.EvaluationInterval(),
		Storage:            ts,
		PrometheusUrl:      web.MustBuildServerUrl(),
	})
	if err := ruleManager.AddRulesFromConfig(conf); err != nil {
		glog.Fatal("Error loading rule files: ", err)
	}
	go ruleManager.Run()

	notificationHandler := notification.NewNotificationHandler(*alertmanagerUrl, notifications)
	go notificationHandler.Run()

	flags := map[string]string{}

	flag.VisitAll(func(f *flag.Flag) {
		flags[f.Name] = f.Value.String()
	})

	prometheusStatus := &web.PrometheusStatusHandler{
		BuildInfo:   BuildInfo,
		Config:      conf.String(),
		RuleManager: ruleManager,
		TargetPools: targetManager.Pools(),
		Flags:       flags,
		Birth:       time.Now(),
	}

	alertsHandler := &web.AlertsHandler{
		RuleManager: ruleManager,
	}

	databasesHandler := &web.DatabasesHandler{
		Provider:        ts.DiskStorage,
		RefreshInterval: 5 * time.Minute,
	}

	metricsService := &api.MetricsService{
		Config:        &conf,
		TargetManager: targetManager,
		Storage:       ts,
	}

	webService := &web.WebService{
		StatusHandler:    prometheusStatus,
		MetricsHandler:   metricsService,
		DatabasesHandler: databasesHandler,
		AlertsHandler:    alertsHandler,
	}

	prometheus := &prometheus{
		compactionTimer: compactionTimer,

		deletionTimer: deletionTimer,

		curationState: prometheusStatus,
		curationSema:  make(chan bool, 1),

		unwrittenSamples: unwrittenSamples,

		stopBackgroundOperations: make(chan bool, 1),

		ruleManager:     ruleManager,
		targetManager:   targetManager,
		notifications:   notifications,
		storage:         ts,
		remoteTSDBQueue: remoteTSDBQueue,
	}
	defer prometheus.close()

	storageStarted := make(chan bool)
	go ts.Serve(storageStarted)
	<-storageStarted

	go prometheus.interruptHandler()

	go func() {
		for _ = range prometheus.compactionTimer.C {
			glog.Info("Starting compaction...")
			err := prometheus.compact(*compactAgeInclusiveness, *compactGroupSize)

			if err != nil {
				glog.Error("could not compact: ", err)
			}
			glog.Info("Done")
		}
	}()

	go func() {
		for _ = range prometheus.deletionTimer.C {
			glog.Info("Starting deletion of stale values...")
			err := prometheus.delete(*deleteAge, deletionBatchSize)

			if err != nil {
				glog.Error("could not delete: ", err)
			}
			glog.Info("Done")
		}
	}()

	go func() {
		err := webService.ServeForever()
		if err != nil {
			glog.Fatal(err)
		}
	}()

	// TODO(all): Migrate this into prometheus.serve().
	for block := range unwrittenSamples {
		if block.Err == nil && len(block.Samples) > 0 {
			ts.AppendSamples(block.Samples)
			if remoteTSDBQueue != nil {
				remoteTSDBQueue.Queue(block.Samples)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"html/template"
	"io/ioutil"
	"regexp"
	"log"

	"net/http"

	"flag"
	//"github.com/FogCreek/mini"

	"github.com/JackKnifed/gnosis"
	"github.com/JackKnifed/blackfriday"
	//"github.com/russross/blackfriday"
)

type WikiPage struct {
	Title string
	//Body  stringo
	Body template.HTML
}

var templates = template.Must(template.ParseFiles("wiki.html"))

var validWiki = regexp.MustCompile("^/([a-zA-Z0-9_ ]+)$")

var	configFile = flag.String("config", "", "specify a configuration file")

//var validFiles = regexp.MustCompile("^/raw/([a-zA-Z0-9]+)\\.(jpg|gif|jpeg|md|png)$")

//var validSupport = regexp.MustCompile("^/global/([a-zA-Z0-9]+)\\.(css|js)$")

func markdownHandler(responsePipe http.ResponseWriter, request *http.Request) {
	validRequest := validWiki.FindStringSubmatch(request.URL.Path)
	config := gnosis.GetConfig()
	if validRequest == nil {
		http.Error(responsePipe, "Request not allowed", 403)
	}
	contents, err := ioutil.ReadFile(config.mainserver.prefix + validRequest[1] + ".md")
	if err != nil {
		http.Error(responsePipe, err.Error(), 500)
	}
	// parse any markdown in the input
	body := template.HTML(blackfriday.MarkdownCommon(contents))

	response := WikiPage{Title: validRequest[1], Body: body}
	err = templates.ExecuteTemplate(responsePipe, "wiki.html", response)
	if err != nil {
		http.Error(responsePipe, err.Error(), http.StatusInternalServerError)
	}
}

/*
func contentHandler(responsePipe http.ResponseWriter, request *http.Request) {
	validRequest := validFiles.FindStringSubmatch(request.URL.Path)
	if validRequest == nil {
		http.Error(responsePipe, "Request not allowed", 403)
	}
	strippedRequest := http.StripPrefix("/raw/", validRequest[0])
	http.ServeFile(responsePipe, request, strippedRequest)
	http.FileServer()
	contents, err := loadFile(validRequest)
	if err != nil {
		http.Error(responsePipe, err, 500)
	}
	response := contents
	err := templates.ExecuteTemplate(responsePipe, "markdown.html", p)
	if err != nil {
		http.Error(responsePipe, err.Error(), http.StatusInternalServerError)
	}

}
*/

/*
func loadFile(filename string) ([]byte, error) {
	contents, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return contents, nil
}
*/

func main() {
	flag.Parse()

	gnosis.LoadConfig(*configFile)

	config := gnosis.GetConfig()

	rawFiles := http.FileServer(http.Dir(config.mainserver.path))
	siteFiles := http.FileServer(http.Dir(config.mainserver.path))

	http.Handle("/raw/", http.StripPrefix("/raw/", rawFiles))
	http.Handle("/site/", http.StripPrefix("/site/", siteFiles))
	http.HandleFunc("/", markdownHandler)

	log.Println(http.ListenAndServe(":"+config.global.port, nil))

<commit_msg>added a } back in that got removed<commit_after>package main

import (
	"html/template"
	"io/ioutil"
	"log"
	"regexp"

	"net/http"

	"flag"
	//"github.com/FogCreek/mini"

	"github.com/JackKnifed/blackfriday"
	"github.com/JackKnifed/gnosis"
	//"github.com/russross/blackfriday"
)

type WikiPage struct {
	Title string
	//Body  stringo
	Body template.HTML
}

var templates = template.Must(template.ParseFiles("wiki.html"))

var validWiki = regexp.MustCompile("^/([a-zA-Z0-9_ ]+)$")

var configFile = flag.String("config", "", "specify a configuration file")

//var validFiles = regexp.MustCompile("^/raw/([a-zA-Z0-9]+)\\.(jpg|gif|jpeg|md|png)$")

//var validSupport = regexp.MustCompile("^/global/([a-zA-Z0-9]+)\\.(css|js)$")

func markdownHandler(responsePipe http.ResponseWriter, request *http.Request) {
	validRequest := validWiki.FindStringSubmatch(request.URL.Path)
	config := gnosis.GetConfig()
	if validRequest == nil {
		http.Error(responsePipe, "Request not allowed", 403)
	}
	contents, err := ioutil.ReadFile(config.mainserver.prefix + validRequest[1] + ".md")
	if err != nil {
		http.Error(responsePipe, err.Error(), 500)
	}
	// parse any markdown in the input
	body := template.HTML(blackfriday.MarkdownCommon(contents))

	response := WikiPage{Title: validRequest[1], Body: body}
	err = templates.ExecuteTemplate(responsePipe, "wiki.html", response)
	if err != nil {
		http.Error(responsePipe, err.Error(), http.StatusInternalServerError)
	}
}

/*
func contentHandler(responsePipe http.ResponseWriter, request *http.Request) {
	validRequest := validFiles.FindStringSubmatch(request.URL.Path)
	if validRequest == nil {
		http.Error(responsePipe, "Request not allowed", 403)
	}
	strippedRequest := http.StripPrefix("/raw/", validRequest[0])
	http.ServeFile(responsePipe, request, strippedRequest)
	http.FileServer()
	contents, err := loadFile(validRequest)
	if err != nil {
		http.Error(responsePipe, err, 500)
	}
	response := contents
	err := templates.ExecuteTemplate(responsePipe, "markdown.html", p)
	if err != nil {
		http.Error(responsePipe, err.Error(), http.StatusInternalServerError)
	}

}
*/

/*
func loadFile(filename string) ([]byte, error) {
	contents, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}
	return contents, nil
}
*/

func main() {
	flag.Parse()

	gnosis.LoadConfig(*configFile)

	config := gnosis.GetConfig()

	rawFiles := http.FileServer(http.Dir(config.mainserver.path))
	siteFiles := http.FileServer(http.Dir(config.mainserver.path))

	http.Handle("/raw/", http.StripPrefix("/raw/", rawFiles))
	http.Handle("/site/", http.StripPrefix("/site/", siteFiles))
	http.HandleFunc("/", markdownHandler)

	log.Println(http.ListenAndServe(":"+config.global.port, nil))

}
<|endoftext|>


<commit_before>// Force12.io is a package that monitors demand for resource in a system and then scales and repurposes
// containers, based on agreed "quality of service" contracts, to best handle that demand within the constraints of your existing VM
// or physical infrastructure (for v1).
//
// Force12 is defined to optimize the use of existing physical and VM resources instantly. VMs cannot be scaled in real time (it takes
// several minutes) and new physical machines take even longer. However, containers can be started or stopped at sub second speeds,
// allowing your infrastructure to adapt itself in real time to meet system demands.
//
// Force12 is aimed at effectively using the resources you have right now - your existing VMs or physical servers - by using them as
// optimally as possible.
//
// The Force12 approach is analogous to the way that a router dynamically optimises the use of a physical network. A router is limited
// by the capacity of the lines physically connected to it. Adding additional capacity is a physical process and takes time. Routers
// therefore make decisions in real time about which packets will be prioritized on a particular line based on the packet's priority
// (defined by a "quality of service" contract).
//
// For example, at times of high bandwidth usage a router might prioritize VOIP traffic over web browsing in real time.
//
// Containers allow Force12 to make similar "instant" judgements on service prioritisation within your existing infrastructure. Routers
// make very simplistic judgments because they have limited time and cpu and they act at a per packet level. Force12 has the capability
// of making far more sophisticated judgements, although even fairly simple ones will still provide a significant new service.
//
// This prototype is a bare bones implementation of Force12.io that recognises only 1 demand type:
// randomised demand for a priority 1 service. Resources are allocated to meet this demand for priority 1, and spare resource can
// be used for a priority 2 service.
//
// These demand type examples have been chosen purely for simplicity of demonstration. In the future more demand types
// will be offered
//
// V1 - Force12.io reacts to increased demand by starting/stopping containers on the slaves already in play.
//
// This version of Force12 starts and stops containers on a Mesos cluser using Marathon as the scheduler
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"net/http"
	"os"
	"time"

	"bitbucket.org/force12io/force12-scheduler/marathon"
	"bitbucket.org/force12io/force12-scheduler/scheduler"
)

type sendStatePayload struct {
	CreatedAt          int64 `json:"createdAt"`
	Priority1Requested int   `json:"priority1Requested"`
	Priority1Running   int   `json:"priority1Running"`
	Priority2Running   int   `json:"priority2Running"`
}

const const_sleep = 100     //milliseconds
const const_stopsleep = 250 //milliseconds pause between stopping and restarting containers
const const_p1demandstart int = 5
const const_p2demandstart int = 4
const const_maxcontainers int = 9

type Demand struct {
	sched scheduler.Scheduler

	p1demand    int // number of Priority 1 tasks demanded
	p2demand    int
	p1requested int // indicates how many P1 tasks we've tried to kick off.
	p2requested int
}

// set returns values that were there (p1, p2)
// if provided value is -1 don't update, demand will always be between 0 and const_maxcontainers
func (d *Demand) set(p1, p2 int) (int, int) {
	//d.mu.Lock()
	p1old := d.p1demand
	p2old := d.p2demand
	if p2 != -1 {
		d.p2demand = p2
	}
	if p1 != -1 {
		d.p1demand = p1
	}
	//d.mu.Unlock()
	return p1old, p2old
}

// get returns client, server AEC - Combine this with the set to reduce code
func (d *Demand) get() (int, int) {
	return d.p1demand, d.p2demand
}

// handle processes a change in demand
// Note that handle will make any judgment on what to do with a demand
// change, including potentially nothing.
func (d *Demand) handle() error {
	var err error
	// AEC NOTE THIS FUNCTION NEEDS TO BE HEAVILY REWRITTEN TO HANDLE ECS
	// WHEN WE PORT THAT OVER TO THE SAME STRUCTURE.
	// THe reason is that all the code we wrote to handle stopping before
	// starting etc.. is handled directly by Marathon so that code
	// from the old scheduler needs to go behind the scheduler interface
	err = d.sched.StopStartNTasks(os.Getenv("CLIENT_TASK"), os.Getenv("CLIENT_FAMILY"), d.p1demand, d.p1requested)
	if err != nil {
		log.Printf("Failed to start Priority1 tasks. %v", err)
	}
	d.sched.StopStartNTasks(os.Getenv("SERVER_TASK"), os.Getenv("SERVER_FAMILY"), d.p2demand, d.p2requested)
	if err != nil {
		log.Printf("Failed to start Priority2 tasks. %v", err)
	}

	return err
}

// update checks for changes in demand, returning true if demand changed
// Note that this function makes no judgement on whether a demand change is
// significant. handle() will determine that.
func (d *Demand) update() bool {
	//log.Println("demand update check.")
	var demandchange bool = false

	container_count, err := d.sched.GetContainerCount("priority1-demand")
	if err != nil {
		log.Printf("Failed to get container count. %v", err)
		return false
	}
	//log.Printf("container count %v\n", container_count)

	//Update our saved p1 demand
	oldP1, _ := d.set(container_count, const_maxcontainers-container_count)

	//Has the demand changed?
	demandchange = (container_count != oldP1)

	if demandchange {
		log.Println("demandchange from, to ", oldP1, container_count)
	}

	return demandchange
}

// sendStateToAPI checks the current state of cluster (or single node) and sends that
// state to the f12 API
func sendStateToAPI(currentdemand *Demand) error {
	count1, count2, err := currentdemand.sched.CountAllTasks()
	if err != nil {
		return fmt.Errorf("Failed to get state err %v", err)
	}

	// Submit a PUT request to the API
	// Note the magic hardcoded string is the user ID, we need to pass this in in some way. ENV VAR?
	url := getBaseF12APIUrl() + "/metrics/" + "5k5gk"
	log.Printf("API PUT: %s", url)

	payload := sendStatePayload{
		CreatedAt:          time.Now().Unix(),
		Priority1Requested: currentdemand.p1demand,
		Priority1Running:   count1,
		Priority2Running:   count2,
	}

	w := &bytes.Buffer{}
	encoder := json.NewEncoder(w)
	err = encoder.Encode(&payload)
	if err != nil {
		return fmt.Errorf("Failed to encode API json. %v", err)
	}

	req, err := http.NewRequest("PUT", url, w)

	if err != nil {
		return fmt.Errorf("Failed to build API PUT request err %v", err)
	}
	//req.Header.Set("X-Custom-Header", "myvalue")
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	defer resp.Body.Close()

	if err != nil {
		// handle error
		return fmt.Errorf("API state err %v", err)
	}

	if resp.StatusCode > 204 {
		return fmt.Errorf("error response from API. %s", resp.Status)
	}
	return err
}

func getBaseF12APIUrl() string {
	baseUrl := os.Getenv("API_ADDRESS")
	if baseUrl == "" {
		baseUrl = "https://force12-windtunnel.herokuapp.com"
	}
	return baseUrl
}

// For the simple prototype, Force12.io sits in a loop checking for demand changes every X milliseconds
// In phase 2 we'll add a reactive mode where appropriate.
//
// Note - we don't route messages from demandcheckers to demandhandlers using channels because we want new values
// to override old values. Queued history is of no importance here.
//
// Also for simplicity this first release is concurrency free (single threaded)
func main() {
	currentdemand := Demand{
		sched: marathon.NewScheduler(),
	}
	currentdemand.set(const_p1demandstart, const_p2demandstart)
	//var errflag bool = false
	var demandchangeflag bool
	//uncomment code below to output logs to file, but there's nothing clever in here to limit file size
	//f, err := os.OpenFile("testlogfile", os.O_RDWR | os.O_CREATE | os.O_APPEND, 0666)
	//if err != nil {
	//   panic(err)
	//}
	//defer f.Close()

	//log.SetOutput(f)
	log.Println("This is a test log entry")

	// Initialise container types
	err := currentdemand.sched.InitScheduler(os.Getenv("CLIENT_TASK"))
	if err != nil {
		log.Printf("Failed to start P1 task. %v", err)
		return
	}
	err = currentdemand.sched.InitScheduler(os.Getenv("SERVER_TASK"))
	if err != nil {
		log.Printf("Failed to start P2 task. %v", err)
		return
	}

	// Find out how many containers we currently have running and get their details
	// Note have decided to do this periodically as a reset as we are getting mysteriously out of whack on ECS
	currentdemand.p1requested, currentdemand.p2requested, err = currentdemand.sched.CountAllTasks()
	if err != nil {
		log.Printf("Failed to count tasks. %v", err)
	}

	//Now we can talk to the DB to check our client demand
	demandchangeflag = currentdemand.update()
	demandchangeflag = true
	var sleepcount float64 = 0

	for {
		switch {
		case demandchangeflag:
			demandchangeflag = false
			//make any changes dictated by this new demand level
			currentdemand.p1requested, currentdemand.p2requested, err = currentdemand.sched.CountAllTasks()
			if err != nil {
				log.Printf("Failed to count tasks. %v", err)
			}
			err = currentdemand.handle()
			if err != nil {
				log.Printf("Failed to handle demand change. %v", err)
			}

		default:
			//log.Println("    .")
			var sleep time.Duration
			sleep = const_sleep * time.Millisecond
			time.Sleep(sleep)
			//Update currentdemand with latest client and server demand, if changed, set flag
			demandchangeflag = currentdemand.update()

			sleepcount++

			//Periodically send state to the API if required
			if os.Getenv("SENDSTATETO_API") == "true" {
				_, frac := math.Modf(math.Mod(sleepcount, 5))
				if frac == 0 {
					err = sendStateToAPI(&currentdemand)
					if err != nil {
						log.Printf("Failed to send state. %v", err)
					}
				}
			}
		}
	}
}
<commit_msg>Simplify the loop<commit_after>// Force12.io is a package that monitors demand for resource in a system and then scales and repurposes
// containers, based on agreed "quality of service" contracts, to best handle that demand within the constraints of your existing VM
// or physical infrastructure (for v1).
//
// Force12 is defined to optimize the use of existing physical and VM resources instantly. VMs cannot be scaled in real time (it takes
// several minutes) and new physical machines take even longer. However, containers can be started or stopped at sub second speeds,
// allowing your infrastructure to adapt itself in real time to meet system demands.
//
// Force12 is aimed at effectively using the resources you have right now - your existing VMs or physical servers - by using them as
// optimally as possible.
//
// The Force12 approach is analogous to the way that a router dynamically optimises the use of a physical network. A router is limited
// by the capacity of the lines physically connected to it. Adding additional capacity is a physical process and takes time. Routers
// therefore make decisions in real time about which packets will be prioritized on a particular line based on the packet's priority
// (defined by a "quality of service" contract).
//
// For example, at times of high bandwidth usage a router might prioritize VOIP traffic over web browsing in real time.
//
// Containers allow Force12 to make similar "instant" judgements on service prioritisation within your existing infrastructure. Routers
// make very simplistic judgments because they have limited time and cpu and they act at a per packet level. Force12 has the capability
// of making far more sophisticated judgements, although even fairly simple ones will still provide a significant new service.
//
// This prototype is a bare bones implementation of Force12.io that recognises only 1 demand type:
// randomised demand for a priority 1 service. Resources are allocated to meet this demand for priority 1, and spare resource can
// be used for a priority 2 service.
//
// These demand type examples have been chosen purely for simplicity of demonstration. In the future more demand types
// will be offered
//
// V1 - Force12.io reacts to increased demand by starting/stopping containers on the slaves already in play.
//
// This version of Force12 starts and stops containers on a Mesos cluser using Marathon as the scheduler
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"net/http"
	"os"
	"time"

	"bitbucket.org/force12io/force12-scheduler/marathon"
	"bitbucket.org/force12io/force12-scheduler/scheduler"
)

type sendStatePayload struct {
	CreatedAt          int64 `json:"createdAt"`
	Priority1Requested int   `json:"priority1Requested"`
	Priority1Running   int   `json:"priority1Running"`
	Priority2Running   int   `json:"priority2Running"`
}

const const_sleep = 100     //milliseconds
const const_stopsleep = 250 //milliseconds pause between stopping and restarting containers
const const_p1demandstart int = 5
const const_p2demandstart int = 4
const const_maxcontainers int = 9

type Demand struct {
	sched scheduler.Scheduler

	p1demand    int // number of Priority 1 tasks demanded
	p2demand    int
	p1requested int // indicates how many P1 tasks we've tried to kick off.
	p2requested int
}

// set returns values that were there (p1, p2)
// if provided value is -1 don't update, demand will always be between 0 and const_maxcontainers
func (d *Demand) set(p1, p2 int) (int, int) {
	//d.mu.Lock()
	p1old := d.p1demand
	p2old := d.p2demand
	if p2 != -1 {
		d.p2demand = p2
	}
	if p1 != -1 {
		d.p1demand = p1
	}
	//d.mu.Unlock()
	return p1old, p2old
}

// get returns client, server AEC - Combine this with the set to reduce code
func (d *Demand) get() (int, int) {
	return d.p1demand, d.p2demand
}

// handle processes a change in demand
// Note that handle will make any judgment on what to do with a demand
// change, including potentially nothing.
func (d *Demand) handle() error {
	var err error
	// AEC NOTE THIS FUNCTION NEEDS TO BE HEAVILY REWRITTEN TO HANDLE ECS
	// WHEN WE PORT THAT OVER TO THE SAME STRUCTURE.
	// THe reason is that all the code we wrote to handle stopping before
	// starting etc.. is handled directly by Marathon so that code
	// from the old scheduler needs to go behind the scheduler interface
	err = d.sched.StopStartNTasks(os.Getenv("CLIENT_TASK"), os.Getenv("CLIENT_FAMILY"), d.p1demand, d.p1requested)
	if err != nil {
		log.Printf("Failed to start Priority1 tasks. %v", err)
	}
	d.sched.StopStartNTasks(os.Getenv("SERVER_TASK"), os.Getenv("SERVER_FAMILY"), d.p2demand, d.p2requested)
	if err != nil {
		log.Printf("Failed to start Priority2 tasks. %v", err)
	}

	return err
}

// update checks for changes in demand, returning true if demand changed
// Note that this function makes no judgement on whether a demand change is
// significant. handle() will determine that.
func (d *Demand) update() bool {
	//log.Println("demand update check.")
	var demandchange bool = false

	container_count, err := d.sched.GetContainerCount("priority1-demand")
	if err != nil {
		log.Printf("Failed to get container count. %v", err)
		return false
	}
	//log.Printf("container count %v\n", container_count)

	//Update our saved p1 demand
	oldP1, _ := d.set(container_count, const_maxcontainers-container_count)

	//Has the demand changed?
	demandchange = (container_count != oldP1)

	if demandchange {
		log.Println("demandchange from, to ", oldP1, container_count)
	}

	return demandchange
}

// sendStateToAPI checks the current state of cluster (or single node) and sends that
// state to the f12 API
func sendStateToAPI(currentdemand *Demand) error {
	count1, count2, err := currentdemand.sched.CountAllTasks()
	if err != nil {
		return fmt.Errorf("Failed to get state err %v", err)
	}

	// Submit a PUT request to the API
	// Note the magic hardcoded string is the user ID, we need to pass this in in some way. ENV VAR?
	url := getBaseF12APIUrl() + "/metrics/" + "5k5gk"
	log.Printf("API PUT: %s", url)

	payload := sendStatePayload{
		CreatedAt:          time.Now().Unix(),
		Priority1Requested: currentdemand.p1demand,
		Priority1Running:   count1,
		Priority2Running:   count2,
	}

	w := &bytes.Buffer{}
	encoder := json.NewEncoder(w)
	err = encoder.Encode(&payload)
	if err != nil {
		return fmt.Errorf("Failed to encode API json. %v", err)
	}

	req, err := http.NewRequest("PUT", url, w)

	if err != nil {
		return fmt.Errorf("Failed to build API PUT request err %v", err)
	}
	//req.Header.Set("X-Custom-Header", "myvalue")
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	defer resp.Body.Close()

	if err != nil {
		// handle error
		return fmt.Errorf("API state err %v", err)
	}

	if resp.StatusCode > 204 {
		return fmt.Errorf("error response from API. %s", resp.Status)
	}
	return err
}

func getBaseF12APIUrl() string {
	baseUrl := os.Getenv("API_ADDRESS")
	if baseUrl == "" {
		baseUrl = "https://force12-windtunnel.herokuapp.com"
	}
	return baseUrl
}

// For the simple prototype, Force12.io sits in a loop checking for demand changes every X milliseconds
// In phase 2 we'll add a reactive mode where appropriate.
//
// Note - we don't route messages from demandcheckers to demandhandlers using channels because we want new values
// to override old values. Queued history is of no importance here.
//
// Also for simplicity this first release is concurrency free (single threaded)
func main() {
	currentdemand := Demand{
		sched: marathon.NewScheduler(),
	}
	currentdemand.set(const_p1demandstart, const_p2demandstart)
	var demandchangeflag bool
	//uncomment code below to output logs to file, but there's nothing clever in here to limit file size
	//f, err := os.OpenFile("testlogfile", os.O_RDWR | os.O_CREATE | os.O_APPEND, 0666)
	//if err != nil {
	//   panic(err)
	//}
	//defer f.Close()

	//log.SetOutput(f)
	log.Println("This is a test log entry")

	// Initialise container types
	err := currentdemand.sched.InitScheduler(os.Getenv("CLIENT_TASK"))
	if err != nil {
		log.Printf("Failed to start P1 task. %v", err)
		return
	}
	err = currentdemand.sched.InitScheduler(os.Getenv("SERVER_TASK"))
	if err != nil {
		log.Printf("Failed to start P2 task. %v", err)
		return
	}

	//Now we can talk to the DB to check our client demand
	demandchangeflag = currentdemand.update()
	demandchangeflag = true
	var sleepcount float64 = 0
	var sleep time.Duration
	sleep = const_sleep * time.Millisecond

	for {
		switch {
		case demandchangeflag:
			demandchangeflag = false
			//make any changes dictated by this new demand level
			currentdemand.p1requested, currentdemand.p2requested, err = currentdemand.sched.CountAllTasks()
			if err != nil {
				log.Printf("Failed to count tasks. %v", err)
			}
			err = currentdemand.handle()
			if err != nil {
				log.Printf("Failed to handle demand change. %v", err)
			}

		default:
			//log.Println("    .")
			time.Sleep(sleep)
			//Update currentdemand with latest client and server demand, if changed, set flag
			demandchangeflag = currentdemand.update()

			sleepcount++

			//Periodically send state to the API if required
			if os.Getenv("SENDSTATETO_API") == "true" {
				_, frac := math.Modf(math.Mod(sleepcount, 5))
				if frac == 0 {
					err = sendStateToAPI(&currentdemand)
					if err != nil {
						log.Printf("Failed to send state. %v", err)
					}
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>// This software is direct fork of https://github.com/knq/chromedp/tree/master/cmd/chromedp-proxy
// with couple of features added
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"

	"github.com/gorilla/websocket"
	"github.com/fatih/color"
)

var (
	flagListen   = flag.String("l", "localhost:9223", "listen address")
	flagRemote   = flag.String("r", "localhost:9222", "remote address")
	flagNoLog    = flag.Bool("n", false, "disable logging to file")
	flagLogMask  = flag.String("log", "logs/cdp-%s.log", "log file mask")
	flagEllipsis = flag.Bool("s", false, "shorten requests and responses")
)

var (
	responseColor = color.New(color.FgHiGreen).SprintfFunc()
	requestColor = color.New(color.FgGreen).SprintfFunc()
	eventsColor = color.New(color.FgHiRed).SprintfFunc()
	protocolColor = color.New(color.FgYellow).SprintfFunc()
	targetColor = color.New(color.FgHiWhite).SprintfFunc()
	methodColor = color.New(color.FgHiYellow).SprintfFunc()
	errorColor = color.New(color.BgRed, color.FgWhite).SprintfFunc()
)

const (
	incomingBufferSize = 10 * 1024 * 1024
	outgoingBufferSize = 25 * 1024 * 1024
	ellipsisLength     = 80
)

var protocolTargetId = center("protocol message", 36)

var wsUpgrader = &websocket.Upgrader{
	ReadBufferSize:  incomingBufferSize,
	WriteBufferSize: outgoingBufferSize,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

var wsDialer = &websocket.Dialer{
	ReadBufferSize:  outgoingBufferSize,
	WriteBufferSize: incomingBufferSize,
}

func main() {
	flag.Parse()

	mux := http.NewServeMux()

	simplep := httputil.NewSingleHostReverseProxy(&url.URL{Scheme: "http", Host: *flagRemote})

	mux.Handle("/json", simplep)
	mux.Handle("/", simplep)

	mux.HandleFunc("/devtools/page/", func(res http.ResponseWriter, req *http.Request) {

		id := path.Base(req.URL.Path)
		f, logger := createLog(id)
		if f != nil {
			defer f.Close()
		}

		stream := make(chan *protocolMessage, 1024)
		go dumpStream(logger, stream)

		logger.Printf(protocolColor("---------- connection from %s ----------", req.RemoteAddr))

		ver, err := checkVersion()
		if err != nil {
			msg := fmt.Sprintf("version error, got: %v", err)
			logger.Println(msg)
			http.Error(res, msg, 500)
			return
		}
		logger.Printf(protocolColor("endpoint %s reported: %s", *flagRemote, serialize(ver)))

		endpoint := "ws://" + *flagRemote + "/devtools/page/" + id

		// connect outgoing websocket
		logger.Printf(protocolColor("connecting to %s", endpoint))
		out, pres, err := wsDialer.Dial(endpoint, nil)
		if err != nil {
			msg := fmt.Sprintf("could not connect to %s, got: %v", endpoint, err)
			logger.Println(protocolColor(msg))
			http.Error(res, msg, 500)
			return
		}
		defer pres.Body.Close()
		defer out.Close()

		logger.Printf(protocolColor("connected to %s", endpoint))

		// connect incoming websocket
		logger.Printf(protocolColor("upgrading connection on %s", req.RemoteAddr))
		in, err := wsUpgrader.Upgrade(res, req, nil)
		if err != nil {
			msg := fmt.Sprintf("could not upgrade websocket from %s, got: %v", req.RemoteAddr, err)
			logger.Println(protocolColor(msg))
			http.Error(res, msg, 500)
			return
		}
		defer in.Close()
		logger.Printf(protocolColor("upgraded connection on %s", req.RemoteAddr))

		ctxt, cancel := context.WithCancel(context.Background())
		defer cancel()

		errc := make(chan error, 1)
		go proxyWS(ctxt, stream, in, out, errc)
		go proxyWS(ctxt, stream, out, in, errc)
		<-errc
		logger.Printf(protocolColor("---------- closing %s ----------", req.RemoteAddr))
	})

	log.Fatal(http.ListenAndServe(*flagListen, mux))
}

func dumpStream(logger *log.Logger, stream chan *protocolMessage) {
	logger.Printf("Legend: %s, %s, %s, %s",
		protocolColor("protocol informations"),
		eventsColor("received events"),
		requestColor("sent requests"),
		responseColor("received responses."),
	)

	requestNames := make(map[uint64]*protocolMessage)
	requestTargetNames := make(map[uint64]*protocolMessage)

	for {
		select {
		case msg := <-stream:
			if msg.InTarget() {
				if msg.IsRequest() {
					requestNames[msg.Id] = nil

					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						requestTargetNames[protocolMessage.Id] = protocolMessage
					} else {
						logger.Printf(protocolColor("Could not deserialize message: %+v", err))
					}
				}

				if msg.IsEvent() {
					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						if protocolMessage.IsEvent() {
							logger.Printf("%s %36s <- %s", targetColor("%s", msg.Params["targetId"]), methodColor(protocolMessage.Method), eventsColor(serialize(protocolMessage.Params)))
						}

						if protocolMessage.IsResponse() {
							if request, ok := requestTargetNames[protocolMessage.Id]; ok {
								if protocolMessage.IsError() {
									logger.Printf("%s %36s(%s) = %s", targetColor("%s", msg.Params["targetId"]), methodColor(request.Method), requestColor(serialize(request.Params)), errorColor(serialize(protocolMessage.Error)))
								} else {
									logger.Printf("%s %36s(%s) = %s", targetColor("%s", msg.Params["targetId"]), methodColor(request.Method), requestColor(serialize(request.Params)), responseColor(serialize(protocolMessage.Result)))
								}
							} else {
								logger.Printf(protocolColor("Could not find target request with id: %d", protocolMessage.Id))
							}
						}
					} else {
						logger.Printf(protocolColor("Could not deserialize message: %+v", err))
					}
				}

			} else {
				if msg.IsRequest() {
					requestNames[msg.Id] = msg
				}

				if msg.IsResponse() {
					if request, ok := requestNames[msg.Id]; ok {
						if request != nil {
							if msg.IsError() {
								logger.Printf("%s %36s(%s) = %s", targetColor(protocolTargetId), methodColor(request.Method), requestColor(serialize(request.Params)), errorColor(serialize(msg.Error)))
							} else {
								logger.Printf("%s %36s(%s) = %s", targetColor(protocolTargetId), methodColor(request.Method), requestColor(serialize(request.Params)), responseColor(serialize(msg.Result)))
							}
						}
					} else {
					logger.Printf(protocolColor("Could not find request with id: %d", msg.Id))
					}
				}

				if msg.IsEvent() {
					logger.Printf("%s %36s <- %s", targetColor(protocolTargetId), methodColor(msg.Method), eventsColor(serialize(msg.Params)))
				}
			}
		}
	}
}

func proxyWS(ctxt context.Context, stream chan *protocolMessage, in, out *websocket.Conn, errc chan error) {
	var mt int
	var buf []byte
	var err error

	for {
		select {
		default:
			mt, buf, err = in.ReadMessage()
			if err != nil {
				errc <- err
				return
			}

			if msg, err := decodeMessage(buf); err == nil {
				stream <- msg
			}

			err = out.WriteMessage(mt, buf)

			if err != nil {
				errc <- err
				return
			}

		case <-ctxt.Done():
			return
		}
	}
}

func checkVersion() (map[string]string, error) {
	cl := &http.Client{}
	req, err := http.NewRequest("GET", "http://" + *flagRemote + "/json/version", nil)
	if err != nil {
		return nil, err
	}

	res, err := cl.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	var v map[string]string
	if err := json.NewDecoder(res.Body).Decode(&v); err != nil {
		return nil, fmt.Errorf("expected json result")
	}

	return v, nil
}

func createLog(id string) (io.Closer, *log.Logger) {
	var f io.Closer
	var w io.Writer = os.Stdout
	if !*flagNoLog && *flagLogMask != "" {
		l, err := os.OpenFile(fmt.Sprintf(*flagLogMask, id), os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
		if err != nil {
			log.Fatal(err)
		}

		f = l
		w = io.MultiWriter(os.Stdout, l)
	}
	return f, log.New(w, "", log.Lmicroseconds)
}

<commit_msg>Protocol errors color<commit_after>// This software is direct fork of https://github.com/knq/chromedp/tree/master/cmd/chromedp-proxy
// with couple of features added
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"path"

	"github.com/gorilla/websocket"
	"github.com/fatih/color"
)

var (
	flagListen   = flag.String("l", "localhost:9223", "listen address")
	flagRemote   = flag.String("r", "localhost:9222", "remote address")
	flagNoLog    = flag.Bool("n", false, "disable logging to file")
	flagLogMask  = flag.String("log", "logs/cdp-%s.log", "log file mask")
	flagEllipsis = flag.Bool("s", false, "shorten requests and responses")
)

var (
	responseColor = color.New(color.FgHiGreen).SprintfFunc()
	requestColor = color.New(color.FgGreen).SprintfFunc()
	eventsColor = color.New(color.FgHiRed).SprintfFunc()
	protocolColor = color.New(color.FgYellow).SprintfFunc()
	protocolError = color.New(color.FgHiYellow, color.BgRed).SprintfFunc()
	targetColor = color.New(color.FgHiWhite).SprintfFunc()
	methodColor = color.New(color.FgHiYellow).SprintfFunc()
	errorColor = color.New(color.BgRed, color.FgWhite).SprintfFunc()
)

const (
	incomingBufferSize = 10 * 1024 * 1024
	outgoingBufferSize = 25 * 1024 * 1024
	ellipsisLength     = 80
)

var protocolTargetId = center("protocol message", 36)

var wsUpgrader = &websocket.Upgrader{
	ReadBufferSize:  incomingBufferSize,
	WriteBufferSize: outgoingBufferSize,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

var wsDialer = &websocket.Dialer{
	ReadBufferSize:  outgoingBufferSize,
	WriteBufferSize: incomingBufferSize,
}

func main() {
	flag.Parse()

	mux := http.NewServeMux()

	simplep := httputil.NewSingleHostReverseProxy(&url.URL{Scheme: "http", Host: *flagRemote})

	mux.Handle("/json", simplep)
	mux.Handle("/", simplep)

	mux.HandleFunc("/devtools/page/", func(res http.ResponseWriter, req *http.Request) {

		id := path.Base(req.URL.Path)
		f, logger := createLog(id)
		if f != nil {
			defer f.Close()
		}

		stream := make(chan *protocolMessage, 1024)
		go dumpStream(logger, stream)

		logger.Printf(protocolColor("---------- connection from %s ----------", req.RemoteAddr))

		ver, err := checkVersion()
		if err != nil {
			msg := fmt.Sprintf("version error, got: %v", err)
			logger.Println(protocolError(msg))
			http.Error(res, msg, 500)
			return
		}
		logger.Printf(protocolColor("endpoint %s reported: %s", *flagRemote, serialize(ver)))

		endpoint := "ws://" + *flagRemote + "/devtools/page/" + id

		// connect outgoing websocket
		logger.Printf(protocolColor("connecting to %s", endpoint))
		out, pres, err := wsDialer.Dial(endpoint, nil)
		if err != nil {
			msg := fmt.Sprintf("could not connect to %s, got: %v", endpoint, err)
			logger.Println(protocolError(msg))
			http.Error(res, msg, 500)
			return
		}
		defer pres.Body.Close()
		defer out.Close()

		logger.Printf(protocolColor("connected to %s", endpoint))

		// connect incoming websocket
		logger.Printf(protocolColor("upgrading connection on %s", req.RemoteAddr))
		in, err := wsUpgrader.Upgrade(res, req, nil)
		if err != nil {
			msg := fmt.Sprintf("could not upgrade websocket from %s, got: %v", req.RemoteAddr, err)
			logger.Println(protocolError(msg))
			http.Error(res, msg, 500)
			return
		}
		defer in.Close()
		logger.Printf(protocolColor("upgraded connection on %s", req.RemoteAddr))

		ctxt, cancel := context.WithCancel(context.Background())
		defer cancel()

		errc := make(chan error, 1)
		go proxyWS(ctxt, stream, in, out, errc)
		go proxyWS(ctxt, stream, out, in, errc)
		<-errc
		logger.Printf(protocolColor("---------- closing %s ----------", req.RemoteAddr))
	})

	log.Fatal(http.ListenAndServe(*flagListen, mux))
}

func dumpStream(logger *log.Logger, stream chan *protocolMessage) {
	logger.Printf("Legend: %s, %s, %s, %s",
		protocolColor("protocol informations"),
		eventsColor("received events"),
		requestColor("sent requests"),
		responseColor("received responses."),
	)

	requestNames := make(map[uint64]*protocolMessage)
	requestTargetNames := make(map[uint64]*protocolMessage)

	for {
		select {
		case msg := <-stream:
			if msg.InTarget() {
				if msg.IsRequest() {
					requestNames[msg.Id] = nil

					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						requestTargetNames[protocolMessage.Id] = protocolMessage
					} else {
						logger.Printf(protocolColor("Could not deserialize message: %+v", err))
					}
				}

				if msg.IsEvent() {
					if protocolMessage, err := decodeMessage([]byte(asString(msg.Params["message"]))); err == nil {
						if protocolMessage.IsEvent() {
							logger.Printf("%s %36s <- %s", targetColor("%s", msg.Params["targetId"]), methodColor(protocolMessage.Method), eventsColor(serialize(protocolMessage.Params)))
						}

						if protocolMessage.IsResponse() {
							if request, ok := requestTargetNames[protocolMessage.Id]; ok {
								if protocolMessage.IsError() {
									logger.Printf("%s %36s(%s) = %s", targetColor("%s", msg.Params["targetId"]), methodColor(request.Method), requestColor(serialize(request.Params)), errorColor(serialize(protocolMessage.Error)))
								} else {
									logger.Printf("%s %36s(%s) = %s", targetColor("%s", msg.Params["targetId"]), methodColor(request.Method), requestColor(serialize(request.Params)), responseColor(serialize(protocolMessage.Result)))
								}
							} else {
								logger.Printf(protocolColor("Could not find target request with id: %d", protocolMessage.Id))
							}
						}
					} else {
						logger.Printf(protocolColor("Could not deserialize message: %+v", err))
					}
				}

			} else {
				if msg.IsRequest() {
					requestNames[msg.Id] = msg
				}

				if msg.IsResponse() {
					if request, ok := requestNames[msg.Id]; ok {
						if request != nil {
							if msg.IsError() {
								logger.Printf("%s %36s(%s) = %s", targetColor(protocolTargetId), methodColor(request.Method), requestColor(serialize(request.Params)), errorColor(serialize(msg.Error)))
							} else {
								logger.Printf("%s %36s(%s) = %s", targetColor(protocolTargetId), methodColor(request.Method), requestColor(serialize(request.Params)), responseColor(serialize(msg.Result)))
							}
						}
					} else {
					logger.Printf(protocolColor("Could not find request with id: %d", msg.Id))
					}
				}

				if msg.IsEvent() {
					logger.Printf("%s %36s <- %s", targetColor(protocolTargetId), methodColor(msg.Method), eventsColor(serialize(msg.Params)))
				}
			}
		}
	}
}

func proxyWS(ctxt context.Context, stream chan *protocolMessage, in, out *websocket.Conn, errc chan error) {
	var mt int
	var buf []byte
	var err error

	for {
		select {
		default:
			mt, buf, err = in.ReadMessage()
			if err != nil {
				errc <- err
				return
			}

			if msg, err := decodeMessage(buf); err == nil {
				stream <- msg
			}

			err = out.WriteMessage(mt, buf)

			if err != nil {
				errc <- err
				return
			}

		case <-ctxt.Done():
			return
		}
	}
}

func checkVersion() (map[string]string, error) {
	cl := &http.Client{}
	req, err := http.NewRequest("GET", "http://" + *flagRemote + "/json/version", nil)
	if err != nil {
		return nil, err
	}

	res, err := cl.Do(req)
	if err != nil {
		return nil, err
	}

	defer res.Body.Close()

	var v map[string]string
	if err := json.NewDecoder(res.Body).Decode(&v); err != nil {
		return nil, fmt.Errorf("expected json result")
	}

	return v, nil
}

func createLog(id string) (io.Closer, *log.Logger) {
	var f io.Closer
	var w io.Writer = os.Stdout
	if !*flagNoLog && *flagLogMask != "" {
		l, err := os.OpenFile(fmt.Sprintf(*flagLogMask, id), os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
		if err != nil {
			log.Fatal(err)
		}

		f = l
		w = io.MultiWriter(os.Stdout, l)
	}
	return f, log.New(w, "", log.Lmicroseconds)
}

<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aerospike/aerospike-client-go"
)

type Config struct {
	Hosts      []string
	Namespace  string
	Sets       []string
	BufferSize int
	Timeout    time.Duration
}

const defaultPort = 3000

func main() {
	c := getConfig()
	client, err := connectToAerospike(c.Hosts)
	if err != nil {
		log.Fatal(err.Error())
	}
	client.DefaultScanPolicy = aerospike.NewScanPolicy()
	client.DefaultScanPolicy.Timeout = c.Timeout

	client.DefaultWritePolicy = aerospike.NewWritePolicy(0, 0)
	client.DefaultWritePolicy.Timeout = c.Timeout

	nodes := client.GetNodes()
	if len(nodes) == 0 {
		log.Fatal("no aerospike nodes found")
	}

	log.Printf("received '%d' nodes, will be using first one to get lists of sets", len(nodes))
	conn, err := nodes[0].GetConnection(c.Timeout)
	if err != nil {
		log.Fatal(err.Error())
	}

	infoMap, err := aerospike.RequestInfo(conn, "sets")
	if err != nil {
		log.Fatal(err.Error())
	}

	sets := extractSets(infoMap["sets"])
	if len(c.Sets) > 0 {
		sets = intersectStrings(sets, c.Sets)
	}

	if len(sets) == 0 {
		log.Println("no sets to clean")
		return
	}

	var wg sync.WaitGroup

	for _, set := range sets {
		wg.Add(1)

		go func(wg *sync.WaitGroup, set string) {
			defer wg.Done()
			log.Printf("starting removing data from set %s", set)
			if err := clearSet(client, c.Namespace, set, c.BufferSize); err != nil {
				log.Printf("clearing set %s failed: %s", set, err.Error())
			}
		}(&wg, set)
	}

	wg.Wait()
}

func getConfig() (c Config) {
	var hostsStr string
	var setsStr string

	flag.StringVar(&hostsStr, "h", "localhost:3000", "Comma-separated list of aerospike hosts")
	flag.StringVar(&c.Namespace, "n", "", "Data namespace")
	flag.StringVar(&setsStr, "s", "", "Comma-separated list of set names to erase")
	flag.IntVar(&c.BufferSize, "b", 25000, "Size of the buffer to pre-load keys before deleting")
	flag.DurationVar(&c.Timeout, "t", 5*time.Second, "Connection timeout")

	flag.Parse()

	c.Hosts = strings.Split(hostsStr, ",")
	if setsStr != "" {
		c.Sets = strings.Split(setsStr, ",")
	}

	return
}

func connectToAerospike(hostStrings []string) (*aerospike.Client, error) {
	if len(hostStrings) == 0 {
		return nil, errors.New("no aerospike host provided")
	}

	hosts := make([]*aerospike.Host, len(hostStrings))
	for i, connStr := range hostStrings {
		hostStr, portStr, err := net.SplitHostPort(connStr)
		if err != nil {
			if isErrNoPort(err) {
				hosts[i] = aerospike.NewHost(connStr, defaultPort)
				continue
			}
			return nil, fmt.Errorf("could not split host string '%s' by host and port: %s", connStr, err.Error())
		}

		port, err := strconv.Atoi(portStr)
		if err != nil {
			return nil, fmt.Errorf("could determine port number from string '%s': %s", connStr, err.Error())
		}
		hosts[i] = aerospike.NewHost(hostStr, port)
	}

	return aerospike.NewClientWithPolicyAndHost(nil, hosts...)
}

func isErrNoPort(err error) bool {
	if addrErr, ok := err.(*net.AddrError); ok {
		return addrErr.Err == "missing port in address"
	}

	return false
}

func extractSets(infoStr string) []string {
	opts := strings.Split(infoStr, ":")
	sets := []string{}
	for _, opt := range opts {
		if len(opt) > 4 && opt[:4] == "set=" {
			sets = append(sets, opt[4:])
		}
	}

	return sets
}

func intersectStrings(sl1, sl2 []string) []string {
	l1 := len(sl1)
	l2 := len(sl2)
	if l1 == 0 || l2 == 0 {
		return []string{}
	}

	minLen := l1
	if l1 > l2 {
		minLen = l2
	}

	res := make([]string, 0, minLen)

	for _, s1 := range sl1 {
		for _, s2 := range sl2 {
			if s1 == s2 {
				res = append(res, s1)
				break
			}
		}
	}

	return res
}

func clearSet(client *aerospike.Client, ns, set string, bufferSize int) error {
	rs, err := client.ScanAll(nil, ns, set)
	if err != nil {
		return err
	}

	buffer := make([]aerospike.Key, 0, bufferSize)
	var wg sync.WaitGroup
	for record := range rs.Results() {
		if record.Err != nil {
			return err
		}

		buffer = append(buffer, *record.Record.Key)

		if len(buffer) == bufferSize {
			wg.Add(1)
			go removeKeys(client, buffer, set, &wg)

			buffer = make([]aerospike.Key, 0, bufferSize)
		}
	}

	if len(buffer) > 0 {
		wg.Add(1)
		go removeKeys(client, buffer, set, &wg)
	}

	wg.Wait()
	return nil
}

func removeKeys(client *aerospike.Client, keys []aerospike.Key, set string, wg *sync.WaitGroup) {
	defer wg.Done()

	log.Printf("removing %d keys in set %s...", len(keys), set)
	removed := 0
	for i := range keys {
		if ok, err := client.Delete(nil, &keys[i]); err != nil {
			log.Printf("removing keys from set %s failed: %s", set, err.Error())
			return
		} else if ok {
			removed++
		}
	}
	log.Printf("removed %d of %d keys in set %s", removed, len(keys), set)
}
<commit_msg>fix sets parse (with legacy support)<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/aerospike/aerospike-client-go"
)

type Config struct {
	Hosts      []string
	Namespace  string
	Sets       []string
	BufferSize int
	Timeout    time.Duration
}

const defaultPort = 3000

func main() {
	c := getConfig()
	client, err := connectToAerospike(c.Hosts)
	if err != nil {
		log.Fatal(err.Error())
	}
	client.DefaultScanPolicy = aerospike.NewScanPolicy()
	client.DefaultScanPolicy.Timeout = c.Timeout

	client.DefaultWritePolicy = aerospike.NewWritePolicy(0, 0)
	client.DefaultWritePolicy.Timeout = c.Timeout

	nodes := client.GetNodes()
	if len(nodes) == 0 {
		log.Fatal("no aerospike nodes found")
	}

	log.Printf("received '%d' nodes, will be using first one to get lists of sets", len(nodes))
	conn, err := nodes[0].GetConnection(c.Timeout)
	if err != nil {
		log.Fatal(err.Error())
	}

	infoMap, err := aerospike.RequestInfo(conn, "sets")
	if err != nil {
		log.Fatal(err.Error())
	}

	sets := extractSets(infoMap["sets"])
	if len(c.Sets) > 0 {
		sets = intersectStrings(sets, c.Sets)
	}

	if len(sets) == 0 {
		log.Println("no sets to clean")
		return
	}

	var wg sync.WaitGroup

	for _, set := range sets {
		wg.Add(1)

		go func(wg *sync.WaitGroup, set string) {
			defer wg.Done()
			log.Printf("starting removing data from set %s", set)
			if err := clearSet(client, c.Namespace, set, c.BufferSize); err != nil {
				log.Printf("clearing set %s failed: %s", set, err.Error())
			}
		}(&wg, set)
	}

	wg.Wait()
}

func getConfig() (c Config) {
	var hostsStr string
	var setsStr string

	flag.StringVar(&hostsStr, "h", "localhost:3000", "Comma-separated list of aerospike hosts")
	flag.StringVar(&c.Namespace, "n", "", "Data namespace")
	flag.StringVar(&setsStr, "s", "", "Comma-separated list of set names to erase")
	flag.IntVar(&c.BufferSize, "b", 25000, "Size of the buffer to pre-load keys before deleting")
	flag.DurationVar(&c.Timeout, "t", 5*time.Second, "Connection timeout")

	flag.Parse()

	c.Hosts = strings.Split(hostsStr, ",")
	if setsStr != "" {
		c.Sets = strings.Split(setsStr, ",")
	}

	return
}

func connectToAerospike(hostStrings []string) (*aerospike.Client, error) {
	if len(hostStrings) == 0 {
		return nil, errors.New("no aerospike host provided")
	}

	hosts := make([]*aerospike.Host, len(hostStrings))
	for i, connStr := range hostStrings {
		hostStr, portStr, err := net.SplitHostPort(connStr)
		if err != nil {
			if isErrNoPort(err) {
				hosts[i] = aerospike.NewHost(connStr, defaultPort)
				continue
			}
			return nil, fmt.Errorf("could not split host string '%s' by host and port: %s", connStr, err.Error())
		}

		port, err := strconv.Atoi(portStr)
		if err != nil {
			return nil, fmt.Errorf("could determine port number from string '%s': %s", connStr, err.Error())
		}
		hosts[i] = aerospike.NewHost(hostStr, port)
	}

	return aerospike.NewClientWithPolicyAndHost(nil, hosts...)
}

func isErrNoPort(err error) bool {
	if addrErr, ok := err.(*net.AddrError); ok {
		return addrErr.Err == "missing port in address"
	}

	return false
}

func extractSets(infoStr string) []string {
	opts := strings.Split(infoStr, ":")
	sets := []string{}
	for _, opt := range opts {
		if len(opt) > 4 && opt[:4] == "set=" {
			// works for aerospike info from v.3.9.0
			sets = append(sets, opt[4:])
		} else if len(opt) > 9 && opt[:9] == "set_name=" {
			// works for legacy versions (below 3.9.0)
			sets = append(sets, opt[9:])
		}
	}

	return sets
}

func intersectStrings(sl1, sl2 []string) []string {
	l1 := len(sl1)
	l2 := len(sl2)
	if l1 == 0 || l2 == 0 {
		return []string{}
	}

	minLen := l1
	if l1 > l2 {
		minLen = l2
	}

	res := make([]string, 0, minLen)

	for _, s1 := range sl1 {
		for _, s2 := range sl2 {
			if s1 == s2 {
				res = append(res, s1)
				break
			}
		}
	}

	return res
}

func clearSet(client *aerospike.Client, ns, set string, bufferSize int) error {
	rs, err := client.ScanAll(nil, ns, set)
	if err != nil {
		return err
	}

	buffer := make([]aerospike.Key, 0, bufferSize)
	var wg sync.WaitGroup
	for record := range rs.Results() {
		if record.Err != nil {
			return err
		}

		buffer = append(buffer, *record.Record.Key)

		if len(buffer) == bufferSize {
			wg.Add(1)
			go removeKeys(client, buffer, set, &wg)

			buffer = make([]aerospike.Key, 0, bufferSize)
		}
	}

	if len(buffer) > 0 {
		wg.Add(1)
		go removeKeys(client, buffer, set, &wg)
	}

	wg.Wait()
	return nil
}

func removeKeys(client *aerospike.Client, keys []aerospike.Key, set string, wg *sync.WaitGroup) {
	defer wg.Done()

	log.Printf("removing %d keys in set %s...", len(keys), set)
	removed := 0
	for i := range keys {
		if ok, err := client.Delete(nil, &keys[i]); err != nil {
			log.Printf("removing keys from set %s failed: %s", set, err.Error())
			return
		} else if ok {
			removed++
		}
	}
	log.Printf("removed %d of %d keys in set %s", removed, len(keys), set)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"fmt"
	"github.com/voxelbrain/goptions"
	"io"
	"os"
	"path"
)

func main() {
	options := struct {
		Htdocs     string `goptions:"--htdocs, description='htdocs directory', obligatory"`
		HtmlOutput string `goptions:"--htmlout, description='HTML output file (relative to htdocs; if \"-\" then stdout will be used)', obligatory"`
		JsOutput   string `goptions:"--jsout, description='JavaScript output file (relative to htdocs)', obligatory"`
		HtmlInput  string `goptions:"--htmlin, description='HTML input file (relative to htdocs; if \"-\" then stdin will be used)', obligatory"`
	}{}

	goptions.ParseAndFail(&options)

	outbuf := &bytes.Buffer{}

	inputStream := io.Reader(os.Stdin)

	if options.HtmlInput != "-" {
		if inf, err := os.Open(path.Join(options.Htdocs, options.HtmlInput)); err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't open input file: %v\n", err)
			os.Exit(1)
		} else {
			inputStream = inf
		}
	}

	tokenizer := html.NewTokenizer(inputStream)

	jsSources := []string{}

	for {
		tokenType := tokenizer.Next()
		if tokenType == html.ErrorToken {
			break
		}

		token := tokenizer.Token()

		if token.Type == html.StartTagToken && token.Data == "script" {
			if src := extractSource(token.Attr); src != "" {
				jsSources = append(jsSources, src)
			}
			continue
		}
		if token.Type == html.EndTagToken && token.Data == "script" {
			// skip closing tag
			continue
		}
		if token.Type == html.SelfClosingTagToken && token.Data == "script" {
			if src := extractSource(token.Attr); src != "" {
				jsSources = append(jsSources, src)
			}
			continue
		}
		if token.Type == html.EndTagToken && token.Data == "html" {
			err := mergeJsSources(options.Htdocs, options.JsOutput, jsSources)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Couldn't merge JavaScript sources: %v", err)
				os.Exit(1)
			}

			scriptStartToken := &html.Token{
				Type: html.StartTagToken,
				Data: "script",
				Attr: []html.Attribute{html.Attribute{Key: "src", Val: options.JsOutput}},
			}
			outbuf.WriteString(scriptStartToken.String())

			scriptEndToken := &html.Token{
				Type: html.EndTagToken,
				Data: "script",
			}
			outbuf.WriteString(scriptEndToken.String())

			linebreakToken := &html.Token{
				Type: html.TextToken,
				Data: "\n",
			}
			outbuf.WriteString(linebreakToken.String())
		}
		outbuf.WriteString(token.String())
	}

	if options.HtmlOutput == "-" {
		os.Stdout.WriteString(outbuf.String())
	} else {
		f, err := os.OpenFile(path.Join(options.Htdocs, options.HtmlOutput), os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't open HTML output file: %v", err)
			os.Exit(1)
		}
		defer f.Close()
		io.Copy(f, outbuf)
	}
}

func extractSource(attr []html.Attribute) string {
	for _, a := range attr {
		if a.Namespace == "" && a.Key == "src" {
			return a.Val
		}
	}
	return ""
}

func mergeJsSources(htdocs, jsOutput string, jsSources []string) error {
	outf, err := os.OpenFile(path.Join(htdocs, jsOutput), os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer outf.Close()

	for _, input := range jsSources {
		inf, err := os.Open(path.Join(htdocs, input))
		if err != nil {
			return err
		}
		defer inf.Close()

		io.Copy(outf, inf)
	}

	return nil
}
<commit_msg>preserve inline script tags but throw away anything inside script tags with src.<commit_after>package main

import (
	"bytes"
	"code.google.com/p/go.net/html"
	"fmt"
	"github.com/voxelbrain/goptions"
	"io"
	"os"
	"path"
)

func main() {
	options := struct {
		Htdocs     string `goptions:"--htdocs, description='htdocs directory', obligatory"`
		HtmlOutput string `goptions:"--htmlout, description='HTML output file (relative to htdocs; if \"-\" then stdout will be used)', obligatory"`
		JsOutput   string `goptions:"--jsout, description='JavaScript output file (relative to htdocs)', obligatory"`
		HtmlInput  string `goptions:"--htmlin, description='HTML input file (relative to htdocs; if \"-\" then stdin will be used)', obligatory"`
	}{}

	goptions.ParseAndFail(&options)

	outbuf := &bytes.Buffer{}

	inputStream := io.Reader(os.Stdin)

	if options.HtmlInput != "-" {
		if inf, err := os.Open(path.Join(options.Htdocs, options.HtmlInput)); err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't open input file: %v\n", err)
			os.Exit(1)
		} else {
			inputStream = inf
		}
	}

	tokenizer := html.NewTokenizer(inputStream)

	jsSources := []string{}

	preserveClosingScriptTag := false
	insideScript := false

	for {
		tokenType := tokenizer.Next()
		if tokenType == html.ErrorToken {
			break
		}

		token := tokenizer.Token()

		if token.Type == html.StartTagToken && token.Data == "script" {
			insideScript = true
			if src, foundSrcAttribute := extractSource(token.Attr); foundSrcAttribute {
				if src != "" {
					jsSources = append(jsSources, src)
				}
			} else {
				outbuf.WriteString(token.String())
				preserveClosingScriptTag = true
			}
			continue
		}
		if token.Type == html.EndTagToken && token.Data == "script" {
			insideScript = false
			if preserveClosingScriptTag {
				outbuf.WriteString(token.String())
				preserveClosingScriptTag = false
			}
			continue
		}
		if token.Type == html.SelfClosingTagToken && token.Data == "script" {
			if src, foundSrcAttribute := extractSource(token.Attr); foundSrcAttribute {
				if src != "" {
					jsSources = append(jsSources, src)
				}
			}
			continue
		}
		if token.Type == html.TextToken && insideScript {
			if preserveClosingScriptTag {
				outbuf.WriteString(token.String())
			}
			continue
		}
		if token.Type == html.EndTagToken && token.Data == "html" {
			err := mergeJsSources(options.Htdocs, options.JsOutput, jsSources)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Couldn't merge JavaScript sources: %v", err)
				os.Exit(1)
			}

			scriptStartToken := &html.Token{
				Type: html.StartTagToken,
				Data: "script",
				Attr: []html.Attribute{html.Attribute{Key: "src", Val: options.JsOutput}},
			}
			outbuf.WriteString(scriptStartToken.String())

			scriptEndToken := &html.Token{
				Type: html.EndTagToken,
				Data: "script",
			}
			outbuf.WriteString(scriptEndToken.String())

			linebreakToken := &html.Token{
				Type: html.TextToken,
				Data: "\n",
			}
			outbuf.WriteString(linebreakToken.String())
		}
		outbuf.WriteString(token.String())
	}

	if options.HtmlOutput == "-" {
		os.Stdout.WriteString(outbuf.String())
	} else {
		f, err := os.OpenFile(path.Join(options.Htdocs, options.HtmlOutput), os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Couldn't open HTML output file: %v", err)
			os.Exit(1)
		}
		defer f.Close()
		io.Copy(f, outbuf)
	}
}

func extractSource(attr []html.Attribute) (src string, foundSrc bool) {
	for _, a := range attr {
		if a.Namespace == "" && a.Key == "src" {
			return a.Val, true
		}
	}
	return "", false
}

func mergeJsSources(htdocs, jsOutput string, jsSources []string) error {
	outf, err := os.OpenFile(path.Join(htdocs, jsOutput), os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer outf.Close()

	for _, input := range jsSources {
		inf, err := os.Open(path.Join(htdocs, input))
		if err != nil {
			return err
		}
		defer inf.Close()

		io.Copy(outf, inf)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"errors"
	"fmt"
	"golang.org/x/net/html/charset"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"regexp"
)

type songInfo struct {
	artist, title string
}

func main() {
	home := os.Getenv("HOME")
	if home == "" {
		log.Fatal("HOME not found")
	}

	dotDir := path.Join(home, ".show-lyrics")
	cacheDir := path.Join(dotDir, "cache")

	for _, dir := range []string{dotDir, cacheDir} {
		err := mkdirUnlessExists(dir)
		if err != nil {
			log.Fatal(err)
		}
	}

	cmusStatus, err := getCmusStatus()
	if err != nil {
		log.Fatal(err)
	}

	songinfo, err := parseCmusStatus(cmusStatus)
	if err != nil {
		log.Fatal(err)
	}

	client := &http.Client{}

	lyrics, err := fetchLyrics(client, songinfo)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(string(lyrics))
}

func mkdirUnlessExists(dir string) error {
	_, err := os.Stat(dir)
	if err != nil {
		err = os.Mkdir(dir, 0755)
		if err != nil {
			return err
		}
	}
	return nil
}

func getCmusStatus() ([]byte, error) {
	cmd := exec.Command("cmus-remote", "-Q")
	return cmd.Output()
}

var artistRe = regexp.MustCompile(`(?m)^tag\s+artist\s+(.+)\s*$`)
var titleRe = regexp.MustCompile(`(?m)^tag\s+title\s+(.+)\s*$`)

func regexpMatch(re *regexp.Regexp, buf []byte) []byte {
	match := re.FindAllSubmatch(buf, 1)
	if len(match) > 0 {
		return match[0][1]
	}
	return nil
}

func parseCmusStatus(cmusStatus []byte) (*songInfo, error) {
	artist := regexpMatch(artistRe, cmusStatus)
	title := regexpMatch(titleRe, cmusStatus)

	if artist == nil || title == nil {
		return nil, errors.New("Failed to parse cmus status")
	}

	si := songInfo{
		artist: string(artist),
		title:  string(title),
	}

	return &si, nil
}

func makeURL(si *songInfo) string {
	artist := []byte(si.artist)
	title := []byte(si.title)

	theRe := regexp.MustCompile(`(?i)^the `)
	weirdRe := regexp.MustCompile(`(?i)[^a-z0-9]`)

	artist = theRe.ReplaceAll(artist, []byte{})

	artist = bytes.ToLower(artist)
	title = bytes.ToLower(title)

	for _, str := range []*[]byte{&artist, &title} {
		*str = bytes.ToLower(*str)
		*str = weirdRe.ReplaceAll(*str, []byte{})
	}

	url := "https://www.azlyrics.com/lyrics/"
	url += string(artist) + "/" + string(title) + ".html"

	return url
}

func fetchLyrics(client *http.Client, si *songInfo) ([]byte, error) {
	reqUrl := makeURL(si)

	req, err := http.NewRequest("GET", reqUrl, nil)
	if err != nil {
		return []byte{}, err
	}

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

	if resp.StatusCode != 200 {
		return []byte{}, errors.New(resp.Status)
	}

	utf8, err := charset.NewReader(resp.Body, resp.Header.Get("Content-Type"))
	if err != nil {
		return []byte{}, err
	}
	body, err := ioutil.ReadAll(utf8)
	if err != nil {
		return []byte{}, err
	}

	lyrics, err := parseLyrics(body)
	if err != nil {
		return []byte{}, err
	}

	return lyrics, nil
}

func htmlStrip(html []byte) []byte {
	commentsRe := regexp.MustCompile(`(?s)<!--.*?-->`)
	brRe := regexp.MustCompile(`<br/?>`)

	html = commentsRe.ReplaceAll(html, []byte{})
	html = brRe.ReplaceAll(html, []byte{})

	// TODO
	// trim leading|trainling whitespace for every line

	return html
}

func parseLyrics(lyricsHtml []byte) ([]byte, error) {
	re := regexp.MustCompile(
		`(?s)<div[^<>]*?class="lyricsh"[^<>]*?>.*?</div>\s*?` +
			`<div[^<>]*?>.*?</div>\s*` +
			`.*?` +
			`<div[^<>]*?>(.*?)</div>`)

	match := re.FindAllSubmatch(lyricsHtml, 1)
	if match == nil {
		return []byte{}, errors.New("Failed to parse html")
	}

	lyrics := htmlStrip(match[0][1])
	return lyrics, nil
}
<commit_msg>Put to cache<commit_after>package main

import (
	"bytes"
	"errors"
	"fmt"
	"golang.org/x/net/html/charset"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"regexp"
)

type songInfo struct {
	artist, title string
}

func main() {
	home := os.Getenv("HOME")
	if home == "" {
		log.Fatal("HOME not found")
	}

	songinfo, err := getSongInfo()
	if err != nil {
		log.Fatal(err)
	}

	// TODO
	// artist and title can contain slashes
	dotDir := path.Join(home, ".show-lyrics")
	cacheDir := path.Join(dotDir, "cache")
	cacheArtistDir := path.Join(cacheDir, songinfo.artist)
	songFile := path.Join(cacheArtistDir, songinfo.title + ".txt")

	for _, dir := range []string{dotDir, cacheDir, cacheArtistDir} {
		err := mkdirUnlessExists(dir)
		if err != nil {
			log.Fatal(err)
		}
	}

	client := &http.Client{}

	lyrics, err := fetchLyrics(client, songinfo)
	if err != nil {
		log.Fatal(err)
	}

	err = ioutil.WriteFile(songFile, lyrics, 0644)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Println(string(lyrics))
}

func getSongInfo() (*songInfo, error) {
	cmusStatus, err := getCmusStatus()
	if err != nil {
		return nil, err
	}
	songinfo, err := parseCmusStatus(cmusStatus)
	if err != nil {
		return nil, err
	}
	return songinfo, nil
}

func mkdirUnlessExists(dir string) error {
	_, err := os.Stat(dir)
	if err != nil {
		err = os.Mkdir(dir, 0755)
		if err != nil {
			return err
		}
	}
	return nil
}

func getCmusStatus() ([]byte, error) {
	cmd := exec.Command("cmus-remote", "-Q")
	return cmd.Output()
}

var artistRe = regexp.MustCompile(`(?m)^tag\s+artist\s+(.+)\s*$`)
var titleRe = regexp.MustCompile(`(?m)^tag\s+title\s+(.+)\s*$`)

func regexpMatch(re *regexp.Regexp, buf []byte) []byte {
	match := re.FindAllSubmatch(buf, 1)
	if len(match) > 0 {
		return match[0][1]
	}
	return nil
}

func parseCmusStatus(cmusStatus []byte) (*songInfo, error) {
	artist := regexpMatch(artistRe, cmusStatus)
	title := regexpMatch(titleRe, cmusStatus)

	if artist == nil || title == nil {
		return nil, errors.New("Failed to parse cmus status")
	}

	si := songInfo{
		artist: string(artist),
		title:  string(title),
	}

	return &si, nil
}

func makeURL(si *songInfo) string {
	artist := []byte(si.artist)
	title := []byte(si.title)

	theRe := regexp.MustCompile(`(?i)^the `)
	weirdRe := regexp.MustCompile(`(?i)[^a-z0-9]`)

	artist = theRe.ReplaceAll(artist, []byte{})

	artist = bytes.ToLower(artist)
	title = bytes.ToLower(title)

	for _, str := range []*[]byte{&artist, &title} {
		*str = bytes.ToLower(*str)
		*str = weirdRe.ReplaceAll(*str, []byte{})
	}

	url := "https://www.azlyrics.com/lyrics/"
	url += string(artist) + "/" + string(title) + ".html"

	return url
}

func fetchLyrics(client *http.Client, si *songInfo) ([]byte, error) {
	reqUrl := makeURL(si)

	req, err := http.NewRequest("GET", reqUrl, nil)
	if err != nil {
		return []byte{}, err
	}

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

	if resp.StatusCode != 200 {
		return []byte{}, errors.New(resp.Status)
	}

	utf8, err := charset.NewReader(resp.Body, resp.Header.Get("Content-Type"))
	if err != nil {
		return []byte{}, err
	}
	body, err := ioutil.ReadAll(utf8)
	if err != nil {
		return []byte{}, err
	}

	lyrics, err := parseLyrics(body)
	if err != nil {
		return []byte{}, err
	}

	return lyrics, nil
}

func htmlStrip(html []byte) []byte {
	commentsRe := regexp.MustCompile(`(?s)<!--.*?-->`)
	brRe := regexp.MustCompile(`<br/?>`)

	html = commentsRe.ReplaceAll(html, []byte{})
	html = brRe.ReplaceAll(html, []byte{})

	// TODO
	// trim leading|trainling whitespace for every line

	return html
}

func parseLyrics(lyricsHtml []byte) ([]byte, error) {
	re := regexp.MustCompile(
		`(?s)<div[^<>]*?class="lyricsh"[^<>]*?>.*?</div>\s*?` +
			`<div[^<>]*?>.*?</div>\s*` +
			`.*?` +
			`<div[^<>]*?>(.*?)</div>`)

	match := re.FindAllSubmatch(lyricsHtml, 1)
	if match == nil {
		return []byte{}, errors.New("Failed to parse html")
	}

	lyrics := htmlStrip(match[0][1])
	return lyrics, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"code.google.com/p/go-netrc/netrc"
	"flag"
	"fmt"
	"log"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

const (
	Version = "0.5"
)

var (
	apiURL = "https://api.heroku.com"
	hkHome = os.Getenv("HOME") + "/.hk"
)

var stdin = bufio.NewReader(os.Stdin)

var updater = Updater{
	url: "https://hk.heroku.com/",
	dir: hkHome + "/update/",
}

type Command struct {
	// args does not include the command name
	Run  func(cmd *Command, args []string)
	Flag flag.FlagSet

	Usage string // first word is the command name
	Short string // `hk help` output
	Long  string // `hk help <cmd>` output
}

func (c *Command) printUsage() {
	fmt.Printf("Usage: hk %s\n\n", c.Usage)
	fmt.Println(strings.TrimSpace(c.Long))
}

func (c *Command) Name() string {
	name := c.Usage
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Running `hk help` will list commands in this order.
var commands = []*Command{
	cmdCreate,
	cmdDestroy,
	cmdCreds,
	cmdEnv,
	cmdFetchUpdate,
	cmdGet,
	cmdSet,
	cmdInfo,
	cmdList,
	cmdOpen,
	cmdPs,
	cmdScale,
	cmdTail,
	cmdRun,
	cmdVersion,
	cmdHelp,
}

var (
	flagApp string // convience var for commands that need it
)

func main() {
	defer updater.run() // doesn't run if os.Exit is called

	log.SetFlags(0)

	args := os.Args[1:]
	if len(args) < 1 {
		usage()
	}

	if s := os.Getenv("HEROKU_API_URL"); s != "" {
		apiURL = strings.TrimRight(s, "/")
	}

	for _, cmd := range commands {
		if cmd.Name() == args[0] {
			cmd.Flag.Usage = usage
			cmd.Flag.Parse(args[1:])
			cmd.Run(cmd, cmd.Flag.Args())
			return
		}
	}

	path := findPlugin(args[0])
	if path == "" {
		fmt.Fprintf(os.Stderr, "Unknown command/plugin: %s\n", args[0])
		usage()
	}
	err := execPlugin(path, args)
	log.Fatal("exec error: ", err)
}

func getCreds(u *url.URL) (user, pass string) {
	if u.User != nil {
		pw, _ := u.User.Password()
		return u.User.Username(), pw
	}

	m, err := netrc.FindMachine(os.Getenv("HOME")+"/.netrc", u.Host)
	if err != nil {
		log.Fatalf("netrc error (%s): %v", u.Host, err)
	}

	return m.Login, m.Password
}

func app() (string, error) {
	if flagApp != "" {
		return flagApp, nil
	}

	b, err := exec.Command("git", "config", "remote.heroku.url").Output()
	if err != nil {
		if isNotFound(err) {
			wdir, _ := os.Getwd()
			return "", fmt.Errorf("could not find git remote heroku in %s", wdir)
		}
		return "", err
	}

	out := strings.Trim(string(b), "\r\n ")

	if !strings.HasPrefix(out, gitURLPre) || !strings.HasSuffix(out, gitURLSuf) {
		return "", fmt.Errorf("could not find app name in heroku git remote")
	}

	// Memoize for later use
	flagApp = out[len(gitURLPre) : len(out)-len(gitURLSuf)]

	return flagApp, nil
}

func isNotFound(err error) bool {
	if ee, ok := err.(*exec.ExitError); ok {
		if ws, ok := ee.ProcessState.Sys().(syscall.WaitStatus); ok {
			return ws.ExitStatus() == 1
		}
	}
	return false
}

func mustApp() string {
	name, err := app()
	if err != nil {
		log.Fatal(err)
	}
	return name
}
<commit_msg>version 0.6<commit_after>package main

import (
	"bufio"
	"code.google.com/p/go-netrc/netrc"
	"flag"
	"fmt"
	"log"
	"net/url"
	"os"
	"os/exec"
	"strings"
	"syscall"
)

const (
	Version = "0.6"
)

var (
	apiURL = "https://api.heroku.com"
	hkHome = os.Getenv("HOME") + "/.hk"
)

var stdin = bufio.NewReader(os.Stdin)

var updater = Updater{
	url: "https://hk.heroku.com/",
	dir: hkHome + "/update/",
}

type Command struct {
	// args does not include the command name
	Run  func(cmd *Command, args []string)
	Flag flag.FlagSet

	Usage string // first word is the command name
	Short string // `hk help` output
	Long  string // `hk help <cmd>` output
}

func (c *Command) printUsage() {
	fmt.Printf("Usage: hk %s\n\n", c.Usage)
	fmt.Println(strings.TrimSpace(c.Long))
}

func (c *Command) Name() string {
	name := c.Usage
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

// Running `hk help` will list commands in this order.
var commands = []*Command{
	cmdCreate,
	cmdDestroy,
	cmdCreds,
	cmdEnv,
	cmdFetchUpdate,
	cmdGet,
	cmdSet,
	cmdInfo,
	cmdList,
	cmdOpen,
	cmdPs,
	cmdScale,
	cmdTail,
	cmdRun,
	cmdVersion,
	cmdHelp,
}

var (
	flagApp string // convience var for commands that need it
)

func main() {
	defer updater.run() // doesn't run if os.Exit is called

	log.SetFlags(0)

	args := os.Args[1:]
	if len(args) < 1 {
		usage()
	}

	if s := os.Getenv("HEROKU_API_URL"); s != "" {
		apiURL = strings.TrimRight(s, "/")
	}

	for _, cmd := range commands {
		if cmd.Name() == args[0] {
			cmd.Flag.Usage = usage
			cmd.Flag.Parse(args[1:])
			cmd.Run(cmd, cmd.Flag.Args())
			return
		}
	}

	path := findPlugin(args[0])
	if path == "" {
		fmt.Fprintf(os.Stderr, "Unknown command/plugin: %s\n", args[0])
		usage()
	}
	err := execPlugin(path, args)
	log.Fatal("exec error: ", err)
}

func getCreds(u *url.URL) (user, pass string) {
	if u.User != nil {
		pw, _ := u.User.Password()
		return u.User.Username(), pw
	}

	m, err := netrc.FindMachine(os.Getenv("HOME")+"/.netrc", u.Host)
	if err != nil {
		log.Fatalf("netrc error (%s): %v", u.Host, err)
	}

	return m.Login, m.Password
}

func app() (string, error) {
	if flagApp != "" {
		return flagApp, nil
	}

	b, err := exec.Command("git", "config", "remote.heroku.url").Output()
	if err != nil {
		if isNotFound(err) {
			wdir, _ := os.Getwd()
			return "", fmt.Errorf("could not find git remote heroku in %s", wdir)
		}
		return "", err
	}

	out := strings.Trim(string(b), "\r\n ")

	if !strings.HasPrefix(out, gitURLPre) || !strings.HasSuffix(out, gitURLSuf) {
		return "", fmt.Errorf("could not find app name in heroku git remote")
	}

	// Memoize for later use
	flagApp = out[len(gitURLPre) : len(out)-len(gitURLSuf)]

	return flagApp, nil
}

func isNotFound(err error) bool {
	if ee, ok := err.(*exec.ExitError); ok {
		if ws, ok := ee.ProcessState.Sys().(syscall.WaitStatus); ok {
			return ws.ExitStatus() == 1
		}
	}
	return false
}

func mustApp() string {
	name, err := app()
	if err != nil {
		log.Fatal(err)
	}
	return name
}
<|endoftext|>