<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
)

func runForward(args []string) error {
	flagset := flag.NewFlagSet("forward", flag.ExitOnError)
	var (
		debug    = flagset.Bool("debug", false, "debug logging")
		apiAddr  = flagset.String("api", "", "listen address for forward API (and metrics)")
		prefixes = stringslice{}
	)
	flagset.Var(&prefixes, "prefix", "prefix annotated on each log record (repeatable)")
	flagset.Usage = usageFor(flagset, "oklog forward [flags] <ingester> [<ingester>...]")
	if err := flagset.Parse(args); err != nil {
		return err
	}
	args = flagset.Args()
	if len(args) <= 0 {
		return errors.New("specify at least one ingest address as an argument")
	}

	// Logging.
	var logger log.Logger
	{
		logLevel := level.AllowInfo()
		if *debug {
			logLevel = level.AllowAll()
		}
		logger = log.NewLogfmtLogger(os.Stderr)
		logger = log.With(logger, "ts", log.DefaultTimestampUTC)
		logger = level.NewFilter(logger, logLevel)
	}

	// Instrumentation.
	forwardBytes := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "oklog",
		Name:      "forward_bytes_total",
		Help:      "Bytes forwarded.",
	})
	forwardRecords := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "oklog",
		Name:      "forward_records_total",
		Help:      "Records forwarded.",
	})
	disconnects := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "oklog",
		Name:      "forward_disconnects",
		Help:      "Number of times forwarder is disconnected from ingester.",
	})
	shortWrites := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "oklog",
		Name:      "forward_short_writes",
		Help:      "Number of times forwarder performs a short write to the ingester.",
	})
	prometheus.MustRegister(
		forwardBytes,
		forwardRecords,
		disconnects,
		shortWrites,
	)

	// For now, just a quick-and-dirty metrics server.
	if *apiAddr != "" {
		apiNetwork, apiAddress, _, _, err := parseAddr(*apiAddr, defaultAPIPort)
		if err != nil {
			return err
		}
		apiListener, err := net.Listen(apiNetwork, apiAddress)
		if err != nil {
			return err
		}
		go func() {
			mux := http.NewServeMux()
			registerMetrics(mux)
			registerProfile(mux)
			panic(http.Serve(apiListener, mux))
		}()
	}

	// Parse URLs for forwarders.
	var urls []*url.URL
	for _, addr := range args {
		schema, host, _, _, err := parseAddr(addr, defaultFastPort)
		if err != nil {
			return errors.Wrap(err, "parsing ingest address")
		}
		u, err := url.Parse(fmt.Sprintf("%s://%s", schema, host))
		if err != nil {
			return errors.Wrap(err, "parsing ingest URL")
		}
		if _, _, err := net.SplitHostPort(u.Host); err != nil {
			return errors.Wrapf(err, "couldn't split host:port")
		}
		urls = append(urls, u)
	}

	// Construct the prefix expression.
	var prefix string
	if len(prefixes) > 0 {
		prefix = strings.Join(prefixes, " ") + " "
	}

	// Shuffle the order.
	rand.Seed(time.Now().UnixNano())
	for i := range urls {
		j := rand.Intn(i + 1)
		urls[i], urls[j] = urls[j], urls[i]
	}

	// Build a scanner for the input, and the last record we scanned.
	// These both outlive any individual connection to an ingester.
	// TODO(pb): have flag for backpressure vs. drop
	var (
		s       = bufio.NewScanner(os.Stdin)
		backoff = time.Duration(0)
	)

	// Enter the connect and forward loop. We do this forever.
	for ; ; urls = append(urls[1:], urls[0]) { // rotate thru URLs
		// We gonna try to connect to this first one.
		target := urls[0]

		host, port, err := net.SplitHostPort(target.Host)
		if err != nil {
			return errors.Wrapf(err, "unexpected error")
		}

		// Support e.g. "tcp+dnssrv://host:port"
		fields := strings.SplitN(target.Scheme, "+", 2)
		if len(fields) == 2 {
			proto, suffix := fields[0], fields[1]
			switch suffix {
			case "dns", "dnsip":
				ips, err := net.LookupIP(host)
				if err != nil {
					level.Warn(logger).Log("LookupIP", host, "err", err)
					backoff = exponential(backoff)
					time.Sleep(backoff)
					continue
				}
				host = ips[rand.Intn(len(ips))].String()
				target.Scheme, target.Host = proto, net.JoinHostPort(host, port)

			case "dnssrv":
				_, records, err := net.LookupSRV("", proto, host)
				if err != nil {
					level.Warn(logger).Log("LookupSRV", host, "err", err)
					backoff = exponential(backoff)
					time.Sleep(backoff)
					continue
				}
				host = records[rand.Intn(len(records))].Target
				target.Scheme, target.Host = proto, net.JoinHostPort(host, port) // TODO(pb): take port from SRV record?

			case "dnsaddr":
				names, err := net.LookupAddr(host)
				if err != nil {
					level.Warn(logger).Log("LookupAddr", host, "err", err)
					backoff = exponential(backoff)
					time.Sleep(backoff)
					continue
				}
				host = names[rand.Intn(len(names))]
				target.Scheme, target.Host = proto, net.JoinHostPort(host, port)

			default:
				level.Warn(logger).Log("unsupported_scheme_suffix", suffix, "using", proto)
				target.Scheme = proto // target.Host stays the same
			}
		}
		level.Debug(logger).Log("raw_target", urls[0].String(), "resolved_target", target.String())

		conn, err := net.Dial(target.Scheme, target.Host)
		if err != nil {
			level.Warn(logger).Log("Dial", target.String(), "err", err)
			backoff = exponential(backoff)
			time.Sleep(backoff)
			continue
		}

		ok := s.Scan()
		for ok {
			// We enter the loop wanting to write s.Text() to the conn.
			record := s.Text()
			if n, err := fmt.Fprintf(conn, "%s%s\n", prefix, record); err != nil {
				disconnects.Inc()
				level.Warn(logger).Log("disconnected_from", target.String(), "due_to", err)
				break
			} else if n < len(record)+1 {
				shortWrites.Inc()
				level.Warn(logger).Log("short_write_to", target.String(), "n", n, "less_than", len(record)+1)
				break // TODO(pb): we should do something more sophisticated here
			}

			// Only once the write succeeds do we scan the next record.
			backoff = 0 // reset the backoff on a successful write
			forwardBytes.Add(float64(len(record)) + 1)
			forwardRecords.Inc()
			ok = s.Scan()
		}
		if !ok {
			level.Info(logger).Log("stdin", "exhausted", "due_to", s.Err())
			return nil
		}
	}
}

func exponential(d time.Duration) time.Duration {
	const (
		min = 16 * time.Millisecond
		max = 1024 * time.Millisecond
	)
	d *= 2
	if d < min {
		d = min
	}
	if d > max {
		d = max
	}
	return d
}
<commit_msg>Judge shortWrite when forward records to ingesters bug (#93)<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
)

func runForward(args []string) error {
	flagset := flag.NewFlagSet("forward", flag.ExitOnError)
	var (
		debug    = flagset.Bool("debug", false, "debug logging")
		apiAddr  = flagset.String("api", "", "listen address for forward API (and metrics)")
		prefixes = stringslice{}
	)
	flagset.Var(&prefixes, "prefix", "prefix annotated on each log record (repeatable)")
	flagset.Usage = usageFor(flagset, "oklog forward [flags] <ingester> [<ingester>...]")
	if err := flagset.Parse(args); err != nil {
		return err
	}
	args = flagset.Args()
	if len(args) <= 0 {
		return errors.New("specify at least one ingest address as an argument")
	}

	// Logging.
	var logger log.Logger
	{
		logLevel := level.AllowInfo()
		if *debug {
			logLevel = level.AllowAll()
		}
		logger = log.NewLogfmtLogger(os.Stderr)
		logger = log.With(logger, "ts", log.DefaultTimestampUTC)
		logger = level.NewFilter(logger, logLevel)
	}

	// Instrumentation.
	forwardBytes := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "oklog",
		Name:      "forward_bytes_total",
		Help:      "Bytes forwarded.",
	})
	forwardRecords := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "oklog",
		Name:      "forward_records_total",
		Help:      "Records forwarded.",
	})
	disconnects := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "oklog",
		Name:      "forward_disconnects",
		Help:      "Number of times forwarder is disconnected from ingester.",
	})
	shortWrites := prometheus.NewCounter(prometheus.CounterOpts{
		Namespace: "oklog",
		Name:      "forward_short_writes",
		Help:      "Number of times forwarder performs a short write to the ingester.",
	})
	prometheus.MustRegister(
		forwardBytes,
		forwardRecords,
		disconnects,
		shortWrites,
	)

	// For now, just a quick-and-dirty metrics server.
	if *apiAddr != "" {
		apiNetwork, apiAddress, _, _, err := parseAddr(*apiAddr, defaultAPIPort)
		if err != nil {
			return err
		}
		apiListener, err := net.Listen(apiNetwork, apiAddress)
		if err != nil {
			return err
		}
		go func() {
			mux := http.NewServeMux()
			registerMetrics(mux)
			registerProfile(mux)
			panic(http.Serve(apiListener, mux))
		}()
	}

	// Parse URLs for forwarders.
	var urls []*url.URL
	for _, addr := range args {
		schema, host, _, _, err := parseAddr(addr, defaultFastPort)
		if err != nil {
			return errors.Wrap(err, "parsing ingest address")
		}
		u, err := url.Parse(fmt.Sprintf("%s://%s", schema, host))
		if err != nil {
			return errors.Wrap(err, "parsing ingest URL")
		}
		if _, _, err := net.SplitHostPort(u.Host); err != nil {
			return errors.Wrapf(err, "couldn't split host:port")
		}
		urls = append(urls, u)
	}

	// Construct the prefix expression.
	var prefix string
	if len(prefixes) > 0 {
		prefix = strings.Join(prefixes, " ") + " "
	}

	// Shuffle the order.
	rand.Seed(time.Now().UnixNano())
	for i := range urls {
		j := rand.Intn(i + 1)
		urls[i], urls[j] = urls[j], urls[i]
	}

	// Build a scanner for the input, and the last record we scanned.
	// These both outlive any individual connection to an ingester.
	// TODO(pb): have flag for backpressure vs. drop
	var (
		s       = bufio.NewScanner(os.Stdin)
		backoff = time.Duration(0)
	)

	// Enter the connect and forward loop. We do this forever.
	for ; ; urls = append(urls[1:], urls[0]) { // rotate thru URLs
		// We gonna try to connect to this first one.
		target := urls[0]

		host, port, err := net.SplitHostPort(target.Host)
		if err != nil {
			return errors.Wrapf(err, "unexpected error")
		}

		// Support e.g. "tcp+dnssrv://host:port"
		fields := strings.SplitN(target.Scheme, "+", 2)
		if len(fields) == 2 {
			proto, suffix := fields[0], fields[1]
			switch suffix {
			case "dns", "dnsip":
				ips, err := net.LookupIP(host)
				if err != nil {
					level.Warn(logger).Log("LookupIP", host, "err", err)
					backoff = exponential(backoff)
					time.Sleep(backoff)
					continue
				}
				host = ips[rand.Intn(len(ips))].String()
				target.Scheme, target.Host = proto, net.JoinHostPort(host, port)

			case "dnssrv":
				_, records, err := net.LookupSRV("", proto, host)
				if err != nil {
					level.Warn(logger).Log("LookupSRV", host, "err", err)
					backoff = exponential(backoff)
					time.Sleep(backoff)
					continue
				}
				host = records[rand.Intn(len(records))].Target
				target.Scheme, target.Host = proto, net.JoinHostPort(host, port) // TODO(pb): take port from SRV record?

			case "dnsaddr":
				names, err := net.LookupAddr(host)
				if err != nil {
					level.Warn(logger).Log("LookupAddr", host, "err", err)
					backoff = exponential(backoff)
					time.Sleep(backoff)
					continue
				}
				host = names[rand.Intn(len(names))]
				target.Scheme, target.Host = proto, net.JoinHostPort(host, port)

			default:
				level.Warn(logger).Log("unsupported_scheme_suffix", suffix, "using", proto)
				target.Scheme = proto // target.Host stays the same
			}
		}
		level.Debug(logger).Log("raw_target", urls[0].String(), "resolved_target", target.String())

		conn, err := net.Dial(target.Scheme, target.Host)
		if err != nil {
			level.Warn(logger).Log("Dial", target.String(), "err", err)
			backoff = exponential(backoff)
			time.Sleep(backoff)
			continue
		}

		ok := s.Scan()
		for ok {
			// We enter the loop wanting to write s.Text() to the conn.
			record := fmt.Sprintf("%s%s\n", prefix, s.Text())
			if n, err := fmt.Fprintf(conn, record); err != nil {
				disconnects.Inc()
				level.Warn(logger).Log("disconnected_from", target.String(), "due_to", err)
				break
			} else if n < len(record) {
				shortWrites.Inc()
				level.Warn(logger).Log("short_write_to", target.String(), "n", n, "less_than", len(record))
				break // TODO(pb): we should do something more sophisticated here
			}

			// Only once the write succeeds do we scan the next record.
			backoff = 0 // reset the backoff on a successful write
			forwardBytes.Add(float64(len(record)))
			forwardRecords.Inc()
			ok = s.Scan()
		}
		if !ok {
			level.Info(logger).Log("stdin", "exhausted", "due_to", s.Err())
			return nil
		}
	}
}

func exponential(d time.Duration) time.Duration {
	const (
		min = 16 * time.Millisecond
		max = 1024 * time.Millisecond
	)
	d *= 2
	if d < min {
		d = min
	}
	if d > max {
		d = max
	}
	return d
}
<|endoftext|>


<commit_before>// Copyright 2021 Google LLC
//
// 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
//
//      https://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 testclient

import (
	"context"
	"fmt"
	"log"
	"strconv"

	"cloud.google.com/go/pubsub"
	"github.com/GoogleCloudPlatform/opentelemetry-operations-e2e-testing/setuptf"
	"google.golang.org/genproto/googleapis/rpc/code"
)

const (
	TestID     string = "test_id"
	Scenario   string = "scenario"
	StatusCode string = "status_code"
	Health     string = "/health"
)

type Request struct {
	// name of the scenario to run
	Scenario string
	TestID   string
	Headers  map[string]string
}

type Response struct {
	StatusCode code.Code
	Headers    map[string]string
}

type Client struct {
	pubsubClient         *pubsub.Client
	requestTopic         *pubsub.Topic
	responseSubscription *pubsub.Subscription
}

func New(ctx context.Context, projectID string, pubsubInfo *setuptf.PubsubInfo) (*Client, error) {
	pubsub, err := pubsub.NewClient(ctx, projectID)
	if err != nil {
		return nil, err
	}
	client := &Client{
		pubsubClient:         pubsub,
		requestTopic:         pubsub.Topic(pubsubInfo.RequestTopic.TopicName),
		responseSubscription: pubsub.Subscription(pubsubInfo.ResponseTopic.SubscriptionName),
	}
	// Disable buffering
	client.requestTopic.PublishSettings.CountThreshold = 1
	return client, nil
}

func (c *Client) Request(
	ctx context.Context,
	request Request,
) (*Response, error) {
	attributes := map[string]string{TestID: request.TestID, "scenario": request.Scenario}
	for k, v := range request.Headers {
		attributes[k] = v
	}
	pubResult := c.requestTopic.Publish(ctx, &pubsub.Message{
		Attributes: attributes,
	})
	messageID, err := pubResult.Get(ctx)
	if err != nil {
		return nil, err
	}

	var (
		res    *Response
		resErr error
	)
	cctx, cancel := context.WithCancel(ctx)
	err = c.responseSubscription.Receive(cctx, func(ctx context.Context, message *pubsub.Message) {
		if testID := message.Attributes[TestID]; testID == request.TestID {
			message.Ack()
			codeInt, err := strconv.Atoi(message.Attributes[StatusCode])
			if err != nil {
				resErr = fmt.Errorf(`response pub/sub message missing required attribute "%v", message: %v`, StatusCode, message)
			} else {
				res = &Response{StatusCode: code.Code(codeInt), Headers: message.Attributes}
			}
			cancel()
		} else {
			message.Nack()
		}
	})

	if err != nil {
		return nil, err
	} else if resErr != nil {
		return nil, resErr
	} else if res == nil {
		// Can happen if cctx times out
		return nil, fmt.Errorf(
			"sent message ID %v, but never received a response on subscription %v",
			messageID,
			c.responseSubscription.String(),
		)
	}

	return res, nil
}

// Call in TestMain() to block until the test server is ready for requests. Uses
// a *log.Logger because this runs before testing.T is available
func (c *Client) WaitForHealth(ctx context.Context, logger *log.Logger) error {
	_, err := c.Request(ctx, Request{Scenario: Health})
	return err
}
<commit_msg>Minor: include error message when converting StatusCode to an integer so it is clear the attribute didn't parse (rather than not being present).<commit_after>// Copyright 2021 Google LLC
//
// 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
//
//      https://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 testclient

import (
	"context"
	"fmt"
	"log"
	"strconv"

	"cloud.google.com/go/pubsub"
	"github.com/GoogleCloudPlatform/opentelemetry-operations-e2e-testing/setuptf"
	"google.golang.org/genproto/googleapis/rpc/code"
)

const (
	TestID     string = "test_id"
	Scenario   string = "scenario"
	StatusCode string = "status_code"
	Health     string = "/health"
)

type Request struct {
	// name of the scenario to run
	Scenario string
	TestID   string
	Headers  map[string]string
}

type Response struct {
	StatusCode code.Code
	Headers    map[string]string
}

type Client struct {
	pubsubClient         *pubsub.Client
	requestTopic         *pubsub.Topic
	responseSubscription *pubsub.Subscription
}

func New(ctx context.Context, projectID string, pubsubInfo *setuptf.PubsubInfo) (*Client, error) {
	pubsub, err := pubsub.NewClient(ctx, projectID)
	if err != nil {
		return nil, err
	}
	client := &Client{
		pubsubClient:         pubsub,
		requestTopic:         pubsub.Topic(pubsubInfo.RequestTopic.TopicName),
		responseSubscription: pubsub.Subscription(pubsubInfo.ResponseTopic.SubscriptionName),
	}
	// Disable buffering
	client.requestTopic.PublishSettings.CountThreshold = 1
	return client, nil
}

func (c *Client) Request(
	ctx context.Context,
	request Request,
) (*Response, error) {
	attributes := map[string]string{TestID: request.TestID, "scenario": request.Scenario}
	for k, v := range request.Headers {
		attributes[k] = v
	}
	pubResult := c.requestTopic.Publish(ctx, &pubsub.Message{
		Attributes: attributes,
	})
	messageID, err := pubResult.Get(ctx)
	if err != nil {
		return nil, err
	}

	var (
		res    *Response
		resErr error
	)
	cctx, cancel := context.WithCancel(ctx)
	err = c.responseSubscription.Receive(cctx, func(ctx context.Context, message *pubsub.Message) {
		if testID := message.Attributes[TestID]; testID == request.TestID {
			message.Ack()
			codeInt, err := strconv.Atoi(message.Attributes[StatusCode])
			if err != nil {
				resErr = fmt.Errorf(`response pub/sub message invalid attribute %q: %v, message: %v`, StatusCode, err, message)
			} else {
				res = &Response{StatusCode: code.Code(codeInt), Headers: message.Attributes}
			}
			cancel()
		} else {
			message.Nack()
		}
	})

	if err != nil {
		return nil, err
	} else if resErr != nil {
		return nil, resErr
	} else if res == nil {
		// Can happen if cctx times out
		return nil, fmt.Errorf(
			"sent message ID %v, but never received a response on subscription %v",
			messageID,
			c.responseSubscription.String(),
		)
	}

	return res, nil
}

// Call in TestMain() to block until the test server is ready for requests. Uses
// a *log.Logger because this runs before testing.T is available
func (c *Client) WaitForHealth(ctx context.Context, logger *log.Logger) error {
	_, err := c.Request(ctx, Request{Scenario: Health})
	return err
}
<|endoftext|>


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

package pipeline

import (
	"fmt"
	"net/http"

	"github.com/spf13/cobra"
	"github.com/spinnaker/spin/cmd/gateclient"

	"github.com/spinnaker/spin/util"
)

type SaveOptions struct {
	*pipelineOptions
	output       string
	pipelineFile string
}

var (
	savePipelineShort = "Save the provided pipeline"
	savePipelineLong  = "Save the provided pipeline"
)

func NewSaveCmd(pipelineOptions pipelineOptions) *cobra.Command {
	options := SaveOptions{
		pipelineOptions: &pipelineOptions,
	}
	cmd := &cobra.Command{
		Use:     "save",
		Aliases: []string{},
		Short:   savePipelineShort,
		Long:    savePipelineLong,
		RunE: func(cmd *cobra.Command, args []string) error {
			return savePipeline(cmd, options)
		},
	}

	cmd.PersistentFlags().StringVarP(&options.pipelineFile, "file", "f", "", "path to the pipeline file")

	return cmd
}

func savePipeline(cmd *cobra.Command, options SaveOptions) error {
	gateClient, err := gateclient.NewGateClient(cmd.InheritedFlags())
	if err != nil {
		return err
	}

	pipelineJson, err := util.ParseJsonFromFileOrStdin(options.pipelineFile, false)
	if err != nil {
		return err
	}
	valid := true
	if _, exists := pipelineJson["name"]; !exists {
		util.UI.Error("Required pipeline key 'name' missing...\n")
		valid = false
	}

	if _, exists := pipelineJson["application"]; !exists {
		util.UI.Error("Required pipeline key 'application' missing...\n")
		valid = false
	}

	if template, exists := pipelineJson["template"]; exists && len(template.(map[string]interface{})) > 0 {
		if _, exists := pipelineJson["schema"]; !exists {
			util.UI.Error("Required pipeline key 'schema' missing for templated pipeline...\n")
			valid = false
		}
	}

	if !valid {
		return fmt.Errorf("Submitted pipeline is invalid: %s\n", pipelineJson)
	}
	pipelineJson["type"] = "templatedPipeline"
	application := pipelineJson["application"].(string)
	pipelineName := pipelineJson["name"].(string)

	foundPipeline, queryResp, _ := gateClient.ApplicationControllerApi.GetPipelineConfigUsingGET(gateClient.Context, application, pipelineName)

	if queryResp.StatusCode != http.StatusOK {
		return fmt.Errorf("Encountered an error querying pipeline, status code: %d\n", queryResp.StatusCode)
	}

	_, exists := pipelineJson["id"].(string)
	var foundPipelineId string
	if len(foundPipeline) > 0 {
		foundPipelineId = foundPipeline["id"].(string)
	}
	if !exists && foundPipelineId != "" {
		pipelineJson["id"] = foundPipelineId
	}

	saveResp, saveErr := gateClient.PipelineControllerApi.SavePipelineUsingPOST(gateClient.Context, pipelineJson)

	if saveErr != nil {
		fmt.Printf("s err: %v", saveErr)
		return err
	}
	if saveResp.StatusCode != http.StatusOK {
		return fmt.Errorf("Encountered an error saving pipeline, status code: %d\n", saveResp.StatusCode)
	}

	util.UI.Info(util.Colorize().Color(fmt.Sprintf("[reset][bold][green]Pipeline save succeeded")))
	return nil
}
<commit_msg>fix(pipeline/save): Fix setting pipeline type. (#129)<commit_after>// Copyright (c) 2018, Google, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.

package pipeline

import (
	"fmt"
	"net/http"

	"github.com/spf13/cobra"
	"github.com/spinnaker/spin/cmd/gateclient"

	"github.com/spinnaker/spin/util"
)

type SaveOptions struct {
	*pipelineOptions
	output       string
	pipelineFile string
}

var (
	savePipelineShort = "Save the provided pipeline"
	savePipelineLong  = "Save the provided pipeline"
)

func NewSaveCmd(pipelineOptions pipelineOptions) *cobra.Command {
	options := SaveOptions{
		pipelineOptions: &pipelineOptions,
	}
	cmd := &cobra.Command{
		Use:     "save",
		Aliases: []string{},
		Short:   savePipelineShort,
		Long:    savePipelineLong,
		RunE: func(cmd *cobra.Command, args []string) error {
			return savePipeline(cmd, options)
		},
	}

	cmd.PersistentFlags().StringVarP(&options.pipelineFile, "file", "f", "", "path to the pipeline file")

	return cmd
}

func savePipeline(cmd *cobra.Command, options SaveOptions) error {
	gateClient, err := gateclient.NewGateClient(cmd.InheritedFlags())
	if err != nil {
		return err
	}

	pipelineJson, err := util.ParseJsonFromFileOrStdin(options.pipelineFile, false)
	if err != nil {
		return err
	}
	valid := true
	if _, exists := pipelineJson["name"]; !exists {
		util.UI.Error("Required pipeline key 'name' missing...\n")
		valid = false
	}

	if _, exists := pipelineJson["application"]; !exists {
		util.UI.Error("Required pipeline key 'application' missing...\n")
		valid = false
	}

	if template, exists := pipelineJson["template"]; exists && len(template.(map[string]interface{})) > 0 {
		if _, exists := pipelineJson["schema"]; !exists {
			util.UI.Error("Required pipeline key 'schema' missing for templated pipeline...\n")
			valid = false
		}
	    pipelineJson["type"] = "templatedPipeline"
	}

	if !valid {
		return fmt.Errorf("Submitted pipeline is invalid: %s\n", pipelineJson)
	}
	application := pipelineJson["application"].(string)
	pipelineName := pipelineJson["name"].(string)

	foundPipeline, queryResp, _ := gateClient.ApplicationControllerApi.GetPipelineConfigUsingGET(gateClient.Context, application, pipelineName)

	if queryResp.StatusCode != http.StatusOK {
		return fmt.Errorf("Encountered an error querying pipeline, status code: %d\n", queryResp.StatusCode)
	}

	_, exists := pipelineJson["id"].(string)
	var foundPipelineId string
	if len(foundPipeline) > 0 {
		foundPipelineId = foundPipeline["id"].(string)
	}
	if !exists && foundPipelineId != "" {
		pipelineJson["id"] = foundPipelineId
	}

	saveResp, saveErr := gateClient.PipelineControllerApi.SavePipelineUsingPOST(gateClient.Context, pipelineJson)

	if saveErr != nil {
		fmt.Printf("s err: %v", saveErr)
		return err
	}
	if saveResp.StatusCode != http.StatusOK {
		return fmt.Errorf("Encountered an error saving pipeline, status code: %d\n", saveResp.StatusCode)
	}

	util.UI.Info(util.Colorize().Color(fmt.Sprintf("[reset][bold][green]Pipeline save succeeded")))
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	log "github.com/marcw/gogol"
	"github.com/marcw/poller"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"time"
)

var (
	configFile = flag.String("config", "", "Configuration file to load checks from")
	httpAddr   = flag.String("hostname", "127.0.0.1:6060", "Adress/host on which the http server will listen.")
	userAgent  = flag.String("ua", "Poller (https://github.com/marcw/poller)", "User agent used by Poller")
	timeout    = flag.Duration("timeout", 10*time.Second, "Timeout")
	backends   = flag.String("backends", "stdout", "Backends to enable. Comma separated.")
	alerts     = flag.String("alerts", "", "Alerts to enable. Comma separated. ie: \"smtp\"")
)

type alertPool map[poller.Alerter]bool
type backendPool map[poller.Backend]bool

func init() {
	flag.Parse()
}

func main() {
	var store poller.Store

	if *configFile != "" {
		store = poller.NewFileStore(*configFile)
	} else {
		store = poller.NewInMemoryStore()
	}
	config := poller.NewConfig(store, poller.NewSimpleScheduler())
	if err := config.Load(); err != nil {
		log.Fatalln(err)
	}

	backendPool, err := instanciateBackendPool()
	if err != nil {
		log.Fatalln(err)
	}
	alerterPool, err := instanciateAlerterPool()
	if err != nil {
		log.Fatalln(err)
	}

	pollerPool := poller.NewHttpPoller(*userAgent, *timeout)
	poller := poller.NewDirectPoller()

	go httpInput(config)
	go poller.Run(config.Scheduler(), backendPool, pollerPool, alerterPool)
	go config.Scheduler().Start()

	select {}
}

func instanciateBackendPool() (*backendPool, error) {
	pool := make(backendPool)
	for _, v := range strings.Split(*backends, ",") {
		switch {
		case v == "statsd":
			statsd, err := poller.NewStatsdBackend()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the Statsd backend: %s", err)
			}
			pool[statsd] = true

		case v == "stdout":
			stdout := poller.NewStdoutBackend()
			pool[stdout] = true

		case v == "librato":
			librato, err := poller.NewLibratoBackend()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the Librato backend: %s", err)
			}
			pool[librato] = true

		case v == "syslog":
			syslog, err := poller.NewSyslogBackend()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the syslog backend: %s", err)
			}
			pool[syslog] = true
		}
	}

	return &pool, nil
}

func instanciateAlerterPool() (*alertPool, error) {
	pool := make(alertPool)
	for _, v := range strings.Split(*alerts, ",") {
		switch {
		case v == "smtp":
			smtp, err := poller.NewSmtpAlerter()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the smtp alerter: %s", err)
			}
			pool[smtp] = true
			break
		case v == "pagerduty":
			smtp, err := poller.NewPagerDutyAlerter()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the pagerduty alerter: %s", err)
			}
			pool[smtp] = true
		}
	}

	return &pool, nil
}

func httpInput(config *poller.Config) {
	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {})
	http.Handle("/checks", poller.NewConfigHttpHandler(config))
	if err := http.ListenAndServe(*httpAddr, nil); err != nil {
		log.Fatalln(err)
	}
}

func (p backendPool) Log(event *poller.Event) {
	for k, _ := range p {
		k.Log(event)
	}
}

func (p backendPool) Close() {
	for k, _ := range p {
		k.Close()
	}
}

func (p alertPool) Alert(event *poller.Event) {
	for k := range p {
		k.Alert(event)
	}
}
<commit_msg>Fix command<commit_after>package main

import (
	"flag"
	"fmt"
	log "github.com/marcw/gogol"
	"github.com/marcw/poller"
	"net/http"
	_ "net/http/pprof"
	"strings"
	"time"
)

var (
	configFile = flag.String("config", "", "Configuration file to load checks from")
	httpAddr   = flag.String("hostname", "127.0.0.1:6060", "Adress/host on which the http server will listen.")
	userAgent  = flag.String("ua", "Poller (https://github.com/marcw/poller)", "User agent used by Poller")
	timeout    = flag.Duration("timeout", 10*time.Second, "Timeout")
	backends   = flag.String("backends", "stdout", "Backends to enable. Comma separated.")
	alerts     = flag.String("alerts", "", "Alerts to enable. Comma separated. ie: \"smtp\"")
)

type alertPool map[poller.Alerter]bool
type backendPool map[poller.Backend]bool

func init() {
	flag.Parse()
}

func main() {
	var store poller.Store

	if *configFile != "" {
		store = poller.NewFileStore(*configFile)
	} else {
		store = poller.NewInMemoryStore()
	}
	config := poller.NewConfig(store, poller.NewSimpleScheduler())
	if err := config.Load(); err != nil {
		log.Fatalln(err)
	}

	backendPool, err := instanciateBackendPool()
	if err != nil {
		log.Fatalln(err)
	}
	alerterPool, err := instanciateAlerterPool()
	if err != nil {
		log.Fatalln(err)
	}

	httpProbe := poller.NewHttpProbe(*userAgent, *timeout)
	poller := poller.NewDirectPoller()

	go httpInput(config)
	go poller.Run(config.Scheduler(), backendPool, httpProbe, alerterPool)
	go config.Scheduler().Start()

	select {}
}

func instanciateBackendPool() (*backendPool, error) {
	pool := make(backendPool)
	for _, v := range strings.Split(*backends, ",") {
		switch {
		case v == "statsd":
			statsd, err := poller.NewStatsdBackend()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the Statsd backend: %s", err)
			}
			pool[statsd] = true

		case v == "stdout":
			stdout := poller.NewStdoutBackend()
			pool[stdout] = true

		case v == "librato":
			librato, err := poller.NewLibratoBackend()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the Librato backend: %s", err)
			}
			pool[librato] = true

		case v == "syslog":
			syslog, err := poller.NewSyslogBackend()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the syslog backend: %s", err)
			}
			pool[syslog] = true
		}
	}

	return &pool, nil
}

func instanciateAlerterPool() (*alertPool, error) {
	pool := make(alertPool)
	for _, v := range strings.Split(*alerts, ",") {
		switch {
		case v == "smtp":
			smtp, err := poller.NewSmtpAlerter()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the smtp alerter: %s", err)
			}
			pool[smtp] = true
			break
		case v == "pagerduty":
			smtp, err := poller.NewPagerDutyAlerter()
			if err != nil {
				return nil, fmt.Errorf("Impossible to instanciate the pagerduty alerter: %s", err)
			}
			pool[smtp] = true
		}
	}

	return &pool, nil
}

func httpInput(config *poller.Config) {
	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {})
	http.Handle("/checks", poller.NewConfigHttpHandler(config))
	if err := http.ListenAndServe(*httpAddr, nil); err != nil {
		log.Fatalln(err)
	}
}

func (p backendPool) Log(event *poller.Event) {
	for k, _ := range p {
		k.Log(event)
	}
}

func (p backendPool) Close() {
	for k, _ := range p {
		k.Close()
	}
}

func (p alertPool) Alert(event *poller.Event) {
	for k := range p {
		k.Alert(event)
	}
}
<|endoftext|>


<commit_before>package main

const Version string = "0.7.6"
<commit_msg>Bump version 0.7.7<commit_after>package main

const Version string = "0.7.7"
<|endoftext|>


<commit_before>package main

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"expvar"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"syscall"
	"time"

	"github.com/ekanite/ekanite"
	"github.com/ekanite/ekanite/input"
	"github.com/ekanite/ekanite/status"
)

var (
	stats = expvar.NewMap("ekanite")
)

// Program parameters
var datadir string
var tcpIface string
var udpIface string
var caPemPath string
var caKeyPath string
var queryIface string
var batchSize int
var batchTimeout int
var indexMaxPending int
var gomaxprocs int
var numShards int
var retentionPeriod string
var cpuProfile string
var memProfile string
var inputFormat string

// Flag set
var fs *flag.FlagSet

// Types
const (
	DefaultDataDir         = "/var/opt/ekanite"
	DefaultBatchSize       = 300
	DefaultBatchTimeout    = 1000
	DefaultIndexMaxPending = 1000
	DefaultNumShards       = 4
	DefaultRetentionPeriod = "168h"
	DefaultQueryAddr       = "localhost:9950"
	DefaultHTTPQueryAddr   = "localhost:8080"
	DefaultDiagsIface      = "localhost:9951"
	DefaultTCPServer       = "localhost:5514"
	DefaultInputFormat     = "syslog"
)

func main() {
	fs = flag.NewFlagSet("", flag.ExitOnError)
	var (
		datadir         = fs.String("datadir", DefaultDataDir, "Set data directory")
		batchSize       = fs.Int("batchsize", DefaultBatchSize, "Indexing batch size")
		batchTimeout    = fs.Int("batchtime", DefaultBatchTimeout, "Indexing batch timeout, in milliseconds")
		indexMaxPending = fs.Int("maxpending", DefaultIndexMaxPending, "Maximum pending index events")
		tcpIface        = fs.String("tcp", DefaultTCPServer, "Syslog server TCP bind address in the form host:port. To disable set to empty string")
		udpIface        = fs.String("udp", "", "Syslog server UDP bind address in the form host:port. If not set, not started")
		diagIface       = fs.String("diag", DefaultDiagsIface, "expvar and pprof bind address in the form host:port. If not set, not started")
		caPemPath       = fs.String("tlspem", "", "path to CA PEM file for TLS-enabled TCP server. If not set, TLS not activated")
		caKeyPath       = fs.String("tlskey", "", "path to CA key file for TLS-enabled TCP server. If not set, TLS not activated")
		queryIface      = fs.String("query", DefaultQueryAddr, "TCP Bind address for query server in the form host:port. To disable set to empty string")
		queryIfaceHttp  = fs.String("queryhttp", DefaultHTTPQueryAddr, "TCP Bind address for http query server in the form host:port. To disable set to empty string")
		numShards       = fs.Int("numshards", DefaultNumShards, "Set number of shards per index")
		retentionPeriod = fs.String("retention", DefaultRetentionPeriod, "Data retention period. Minimum is 24 hours")
		cpuProfile      = fs.String("cpuprof", "", "Where to write CPU profiling data. Not written if not set")
		memProfile      = fs.String("memprof", "", "Where to write memory profiling data. Not written if not set")
		inputFormat     = fs.String("input", DefaultInputFormat, "Message format of input (only syslog supported)")
	)
	fs.Usage = printHelp
	fs.Parse(os.Args[1:])

	absDataDir, err := filepath.Abs(*datadir)
	if err != nil {
		log.Fatalf("failed to get absolute data path for '%s': %s", *datadir, err.Error())
	}

	// Get the retention period.
	retention, err := time.ParseDuration(*retentionPeriod)
	if err != nil {
		log.Fatalf("failed to parse retention period '%s'", *retentionPeriod)
	}

	log.SetFlags(log.LstdFlags)
	log.SetPrefix("[ekanite] ")
	log.Printf("ekanite started using %s for index storage", absDataDir)

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("GOMAXPROCS set to", runtime.GOMAXPROCS(0))

	// Start the expvar handler if requested.
	if *diagIface != "" {
		startDiagServer(*diagIface)
	}

	// Create and open the Engine.
	engine := ekanite.NewEngine(absDataDir)
	if engine == nil {
		log.Fatalf("failed to create indexing engine at %s", absDataDir)
	}
	engine.NumShards = *numShards
	engine.RetentionPeriod = retention

	if err := engine.Open(); err != nil {
		log.Fatalf("failed to open engine: %s", err.Error())
	}
	log.Printf("engine opened with shard number of %d, retention period of %s",
		engine.NumShards, engine.RetentionPeriod)

	// Start the simple query server if requested.
	if *queryIface != "" {
		startQueryServer(*queryIface, engine)
	}

	// Start the http query server if requested.
	if *queryIfaceHttp != "" {
		startHTTPQueryServer(*queryIfaceHttp, engine)
	}

	// Create and start the batcher.
	batcherTimeout := time.Duration(*batchTimeout) * time.Millisecond
	batcher := ekanite.NewBatcher(engine, *batchSize, batcherTimeout, *indexMaxPending)
	if batcher == nil {
		log.Fatal("failed to create indexing batcher")
	}

	errChan := make(chan error)
	if err := batcher.Start(errChan); err != nil {
		log.Fatalf("failed to start indexing batcher: %s", err.Error())
	}
	log.Printf("batching configured with size %d, timeout %s, max pending %d",
		*batchSize, batcherTimeout, *indexMaxPending)

	// Start draining batcher errors.
	go drainLog("error indexing batch", errChan)

	// Start TCP collector if requested.
	if *tcpIface != "" {
		var tlsConfig *tls.Config
		if *caPemPath != "" && *caKeyPath != "" {
			tlsConfig, err = newTLSConfig(*caPemPath, *caKeyPath)
			if err != nil {
				log.Fatalf("failed to configure TLS: %s", err.Error())
			}
			log.Printf("TLS successfully configured")
		}

		collector, err := input.NewCollector("tcp", *tcpIface, *inputFormat, tlsConfig)
		if err != nil {
			log.Fatalf("failed to create TCP collector: %s", err.Error())
		}
		if err := collector.Start(batcher.C()); err != nil {
			log.Fatalf("failed to start TCP collector: %s", err.Error())
		}
		log.Printf("TCP collector listening to %s", *tcpIface)
	}

	// Start UDP collector if requested.
	if *udpIface != "" {
		if err := startUDPCollector(*udpIface, *inputFormat, batcher); err != nil {
			log.Fatalf("failed to start UDP collector: %s", err.Error())
		}
		log.Printf("UDP collector listening to %s", *udpIface)
	}

	// Start profiling.
	startProfile(*cpuProfile, *memProfile)

	stats.Set("launch", time.Now().UTC())

	// Set up signal handling.
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Block until one of the signals above is received
	select {
	case <-signalCh:
		log.Println("signal received, shutting down...")
	}

	stopProfile()
}

func startTCPCollector(iface, format string, tls *tls.Config, batcher *ekanite.Batcher) error {
    return nil
}

func startUDPCollector(iface, format string, batcher *ekanite.Batcher) error {
	collector, err := input.NewCollector("udp", iface, format, nil)
	if err != nil {
		return fmt.Errorf("failed to create UDP collector: %s", err.Error())
	}
	if err := collector.Start(batcher.C()); err != nil {
		return fmt.Errorf("failed to start UDP collector: %s", err.Error())
	}

	return nil
}

func startQueryServer(iface string, engine *ekanite.Engine) {
	server := ekanite.NewServer(iface, engine)
	if server == nil {
		log.Fatal("failed to create query server")
	}
	if err := server.Start(); err != nil {
		log.Fatalf("failed to start query server: %s", err.Error())
	}
	log.Printf("query server listening on %s", iface)
}

func startHTTPQueryServer(iface string, engine *ekanite.Engine) {
	server := ekanite.NewHTTPServer(iface, engine)
	if server == nil {
		log.Fatal("failed to create HTTP query server")
	}
	if err := server.Start(); err != nil {
		log.Fatalf("failed to start HTTP query server: %s", err.Error())
	}
	log.Printf("HTTP query server listening on %s", iface)
}

func startDiagServer(iface string) {
	diagServer := status.NewService(iface)
	if err := diagServer.Start(); err != nil {
		log.Fatalf("failed to start status server on %s: %s", iface, err.Error())
	}
	log.Printf("diagnostic server listening on %s", iface)
}

func newTLSConfig(caPemPath, caKeyPath string) (*tls.Config, error) {
	var config *tls.Config

	caPem, err := ioutil.ReadFile(caPemPath)
	if err != nil {
		return nil, err
	}
	ca, err := x509.ParseCertificate(caPem)
	if err != nil {
		return nil, err
	}

	caKey, err := ioutil.ReadFile(caKeyPath)
	if err != nil {
		return nil, err
	}
	key, err := x509.ParsePKCS1PrivateKey(caKey)
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	pool.AddCert(ca)

	cert := tls.Certificate{
		Certificate: [][]byte{caPem},
		PrivateKey:  key,
	}

	config = &tls.Config{
		ClientAuth:   tls.RequireAndVerifyClientCert,
		Certificates: []tls.Certificate{cert},
		ClientCAs:    pool,
	}

	config.Rand = rand.Reader

	return config, nil
}

// drainLog drains errors from the channel and simply logs them
func drainLog(msg string, errChan <-chan error) {
	for {
		select {
		case err := <-errChan:
			if err != nil {
				log.Printf("%s: %s", msg, err.Error())
			}
		}
	}
}

// prof stores the file locations of active profiles.
var prof struct {
	cpu *os.File
	mem *os.File
}

// StartProfile initializes the cpu and memory profile, if specified.
func startProfile(cpuprofile, memprofile string) {
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatalf("cpuprofile: %v", err)
		}
		log.Printf("writing CPU profile to: %s\n", cpuprofile)
		prof.cpu = f
		pprof.StartCPUProfile(prof.cpu)
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatalf("memprofile: %v", err)
		}
		log.Printf("writing memory profile to: %s\n", memprofile)
		prof.mem = f
		runtime.MemProfileRate = 4096
	}

}

// StopProfile closes the cpu and memory profiles if they are running.
func stopProfile() {
	if prof.cpu != nil {
		pprof.StopCPUProfile()
		prof.cpu.Close()
		log.Println("CPU profile stopped")
	}
	if prof.mem != nil {
		pprof.Lookup("heap").WriteTo(prof.mem, 0)
		prof.mem.Close()
		log.Println("memory profile stopped")
	}
}

func printHelp() {
	fmt.Println("ekanited [options]")
	fs.PrintDefaults()
}
<commit_msg>Complete TCP start refactor<commit_after>package main

import (
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"expvar"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"syscall"
	"time"

	"github.com/ekanite/ekanite"
	"github.com/ekanite/ekanite/input"
	"github.com/ekanite/ekanite/status"
)

var (
	stats = expvar.NewMap("ekanite")
)

// Program parameters
var datadir string
var tcpIface string
var udpIface string
var caPemPath string
var caKeyPath string
var queryIface string
var batchSize int
var batchTimeout int
var indexMaxPending int
var gomaxprocs int
var numShards int
var retentionPeriod string
var cpuProfile string
var memProfile string
var inputFormat string

// Flag set
var fs *flag.FlagSet

// Types
const (
	DefaultDataDir         = "/var/opt/ekanite"
	DefaultBatchSize       = 300
	DefaultBatchTimeout    = 1000
	DefaultIndexMaxPending = 1000
	DefaultNumShards       = 4
	DefaultRetentionPeriod = "168h"
	DefaultQueryAddr       = "localhost:9950"
	DefaultHTTPQueryAddr   = "localhost:8080"
	DefaultDiagsIface      = "localhost:9951"
	DefaultTCPServer       = "localhost:5514"
	DefaultInputFormat     = "syslog"
)

func main() {
	fs = flag.NewFlagSet("", flag.ExitOnError)
	var (
		datadir         = fs.String("datadir", DefaultDataDir, "Set data directory")
		batchSize       = fs.Int("batchsize", DefaultBatchSize, "Indexing batch size")
		batchTimeout    = fs.Int("batchtime", DefaultBatchTimeout, "Indexing batch timeout, in milliseconds")
		indexMaxPending = fs.Int("maxpending", DefaultIndexMaxPending, "Maximum pending index events")
		tcpIface        = fs.String("tcp", DefaultTCPServer, "Syslog server TCP bind address in the form host:port. To disable set to empty string")
		udpIface        = fs.String("udp", "", "Syslog server UDP bind address in the form host:port. If not set, not started")
		diagIface       = fs.String("diag", DefaultDiagsIface, "expvar and pprof bind address in the form host:port. If not set, not started")
		caPemPath       = fs.String("tlspem", "", "path to CA PEM file for TLS-enabled TCP server. If not set, TLS not activated")
		caKeyPath       = fs.String("tlskey", "", "path to CA key file for TLS-enabled TCP server. If not set, TLS not activated")
		queryIface      = fs.String("query", DefaultQueryAddr, "TCP Bind address for query server in the form host:port. To disable set to empty string")
		queryIfaceHttp  = fs.String("queryhttp", DefaultHTTPQueryAddr, "TCP Bind address for http query server in the form host:port. To disable set to empty string")
		numShards       = fs.Int("numshards", DefaultNumShards, "Set number of shards per index")
		retentionPeriod = fs.String("retention", DefaultRetentionPeriod, "Data retention period. Minimum is 24 hours")
		cpuProfile      = fs.String("cpuprof", "", "Where to write CPU profiling data. Not written if not set")
		memProfile      = fs.String("memprof", "", "Where to write memory profiling data. Not written if not set")
		inputFormat     = fs.String("input", DefaultInputFormat, "Message format of input (only syslog supported)")
	)
	fs.Usage = printHelp
	fs.Parse(os.Args[1:])

	absDataDir, err := filepath.Abs(*datadir)
	if err != nil {
		log.Fatalf("failed to get absolute data path for '%s': %s", *datadir, err.Error())
	}

	// Get the retention period.
	retention, err := time.ParseDuration(*retentionPeriod)
	if err != nil {
		log.Fatalf("failed to parse retention period '%s'", *retentionPeriod)
	}

	log.SetFlags(log.LstdFlags)
	log.SetPrefix("[ekanite] ")
	log.Printf("ekanite started using %s for index storage", absDataDir)

	runtime.GOMAXPROCS(runtime.NumCPU())
	log.Println("GOMAXPROCS set to", runtime.GOMAXPROCS(0))

	// Start the expvar handler if requested.
	if *diagIface != "" {
		startDiagServer(*diagIface)
	}

	// Create and open the Engine.
	engine := ekanite.NewEngine(absDataDir)
	if engine == nil {
		log.Fatalf("failed to create indexing engine at %s", absDataDir)
	}
	engine.NumShards = *numShards
	engine.RetentionPeriod = retention

	if err := engine.Open(); err != nil {
		log.Fatalf("failed to open engine: %s", err.Error())
	}
	log.Printf("engine opened with shard number of %d, retention period of %s",
		engine.NumShards, engine.RetentionPeriod)

	// Start the simple query server if requested.
	if *queryIface != "" {
		startQueryServer(*queryIface, engine)
	}

	// Start the http query server if requested.
	if *queryIfaceHttp != "" {
		startHTTPQueryServer(*queryIfaceHttp, engine)
	}

	// Create and start the batcher.
	batcherTimeout := time.Duration(*batchTimeout) * time.Millisecond
	batcher := ekanite.NewBatcher(engine, *batchSize, batcherTimeout, *indexMaxPending)
	if batcher == nil {
		log.Fatal("failed to create indexing batcher")
	}

	errChan := make(chan error)
	if err := batcher.Start(errChan); err != nil {
		log.Fatalf("failed to start indexing batcher: %s", err.Error())
	}
	log.Printf("batching configured with size %d, timeout %s, max pending %d",
		*batchSize, batcherTimeout, *indexMaxPending)

	// Start draining batcher errors.
	go drainLog("error indexing batch", errChan)

	// Start TCP collector if requested.
	if *tcpIface != "" {
		var tlsConfig *tls.Config
		if *caPemPath != "" && *caKeyPath != "" {
			tlsConfig, err = newTLSConfig(*caPemPath, *caKeyPath)
			if err != nil {
				log.Fatalf("failed to configure TLS: %s", err.Error())
			}
			log.Printf("TLS successfully configured")
		}

		if err := startTCPCollector(*tcpIface, *inputFormat, tlsConfig, batcher); err != nil {
			log.Fatalf("failed to start TCP collector: %s", err.Error())
		}
		log.Printf("TCP collector listening to %s", *tcpIface)
	}

	// Start UDP collector if requested.
	if *udpIface != "" {
		if err := startUDPCollector(*udpIface, *inputFormat, batcher); err != nil {
			log.Fatalf("failed to start UDP collector: %s", err.Error())
		}
		log.Printf("UDP collector listening to %s", *udpIface)
	}

	// Start profiling.
	startProfile(*cpuProfile, *memProfile)

	stats.Set("launch", time.Now().UTC())

	// Set up signal handling.
	signalCh := make(chan os.Signal, 1)
	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)

	// Block until one of the signals above is received
	select {
	case <-signalCh:
		log.Println("signal received, shutting down...")
	}

	stopProfile()
}

func startTCPCollector(iface, format string, tls *tls.Config, batcher *ekanite.Batcher) error {
	collector, err := input.NewCollector("tcp", iface, format, tls)
	if err != nil {
		return fmt.Errorf(("failed to create TCP collector: %s"), err.Error())
	}
	if err := collector.Start(batcher.C()); err != nil {
		return fmt.Errorf("failed to start TCP collector: %s", err.Error())
	}

	return nil
}

func startUDPCollector(iface, format string, batcher *ekanite.Batcher) error {
	collector, err := input.NewCollector("udp", iface, format, nil)
	if err != nil {
		return fmt.Errorf("failed to create UDP collector: %s", err.Error())
	}
	if err := collector.Start(batcher.C()); err != nil {
		return fmt.Errorf("failed to start UDP collector: %s", err.Error())
	}

	return nil
}

func startQueryServer(iface string, engine *ekanite.Engine) {
	server := ekanite.NewServer(iface, engine)
	if server == nil {
		log.Fatal("failed to create query server")
	}
	if err := server.Start(); err != nil {
		log.Fatalf("failed to start query server: %s", err.Error())
	}
	log.Printf("query server listening on %s", iface)
}

func startHTTPQueryServer(iface string, engine *ekanite.Engine) {
	server := ekanite.NewHTTPServer(iface, engine)
	if server == nil {
		log.Fatal("failed to create HTTP query server")
	}
	if err := server.Start(); err != nil {
		log.Fatalf("failed to start HTTP query server: %s", err.Error())
	}
	log.Printf("HTTP query server listening on %s", iface)
}

func startDiagServer(iface string) {
	diagServer := status.NewService(iface)
	if err := diagServer.Start(); err != nil {
		log.Fatalf("failed to start status server on %s: %s", iface, err.Error())
	}
	log.Printf("diagnostic server listening on %s", iface)
}

func newTLSConfig(caPemPath, caKeyPath string) (*tls.Config, error) {
	var config *tls.Config

	caPem, err := ioutil.ReadFile(caPemPath)
	if err != nil {
		return nil, err
	}
	ca, err := x509.ParseCertificate(caPem)
	if err != nil {
		return nil, err
	}

	caKey, err := ioutil.ReadFile(caKeyPath)
	if err != nil {
		return nil, err
	}
	key, err := x509.ParsePKCS1PrivateKey(caKey)
	if err != nil {
		return nil, err
	}
	pool := x509.NewCertPool()
	pool.AddCert(ca)

	cert := tls.Certificate{
		Certificate: [][]byte{caPem},
		PrivateKey:  key,
	}

	config = &tls.Config{
		ClientAuth:   tls.RequireAndVerifyClientCert,
		Certificates: []tls.Certificate{cert},
		ClientCAs:    pool,
	}

	config.Rand = rand.Reader

	return config, nil
}

// drainLog drains errors from the channel and simply logs them
func drainLog(msg string, errChan <-chan error) {
	for {
		select {
		case err := <-errChan:
			if err != nil {
				log.Printf("%s: %s", msg, err.Error())
			}
		}
	}
}

// prof stores the file locations of active profiles.
var prof struct {
	cpu *os.File
	mem *os.File
}

// StartProfile initializes the cpu and memory profile, if specified.
func startProfile(cpuprofile, memprofile string) {
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatalf("cpuprofile: %v", err)
		}
		log.Printf("writing CPU profile to: %s\n", cpuprofile)
		prof.cpu = f
		pprof.StartCPUProfile(prof.cpu)
	}

	if memprofile != "" {
		f, err := os.Create(memprofile)
		if err != nil {
			log.Fatalf("memprofile: %v", err)
		}
		log.Printf("writing memory profile to: %s\n", memprofile)
		prof.mem = f
		runtime.MemProfileRate = 4096
	}

}

// StopProfile closes the cpu and memory profiles if they are running.
func stopProfile() {
	if prof.cpu != nil {
		pprof.StopCPUProfile()
		prof.cpu.Close()
		log.Println("CPU profile stopped")
	}
	if prof.mem != nil {
		pprof.Lookup("heap").WriteTo(prof.mem, 0)
		prof.mem.Close()
		log.Println("memory profile stopped")
	}
}

func printHelp() {
	fmt.Println("ekanited [options]")
	fs.PrintDefaults()
}
<|endoftext|>


<commit_before>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
The gomote command is a client for the Go builder infrastructure.
It's a remote control for remote Go builder machines.

Usage:

  gomote [global-flags] cmd [cmd-flags]

  For example,
  $ gomote create openbsd-amd64-gce56
  $ gomote push
  $ gomote run openbsd-amd64-gce56 src/make.bash

TODO: document more, and figure out the CLI interface more.
*/
package main

import (
	"flag"
	"fmt"
	"os"
	"sort"

	"golang.org/x/build/buildenv"
	"golang.org/x/build/buildlet"
)

var (
	buildEnv *buildenv.Environment
)

type command struct {
	name string
	des  string
	run  func([]string) error
}

var commands = map[string]command{}

func sortedCommands() []string {
	s := make([]string, 0, len(commands))
	for name := range commands {
		s = append(s, name)
	}
	sort.Strings(s)
	return s
}

func usage() {
	fmt.Fprintf(os.Stderr, `Usage of gomote: gomote [global-flags] <cmd> [cmd-flags]

Global flags:
`)
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "Commands:\n\n")
	for _, name := range sortedCommands() {
		fmt.Fprintf(os.Stderr, "  %-10s %s\n", name, commands[name].des)
	}
	os.Exit(1)
}

func registerCommand(name, des string, run func([]string) error) {
	if _, dup := commands[name]; dup {
		panic("duplicate registration of " + name)
	}
	commands[name] = command{
		name: name,
		des:  des,
		run:  run,
	}
}

func registerCommands() {
	registerCommand("create", "create a buildlet", create)
	registerCommand("destroy", "destroy a buildlet", destroy)
	registerCommand("gettar", "extract a tar.gz from a buildlet", getTar)
	registerCommand("ls", "list the contents of a directory on a buildlet", ls)
	registerCommand("list", "list buildlets", list)
	registerCommand("ping", "test whether a buildlet is alive and reachable ", ping)
	registerCommand("push", "sync the repo of your pwd to the buildlet", push)
	registerCommand("put", "put files on a buildlet", put)
	registerCommand("put14", "put Go 1.4 in place", put14)
	registerCommand("puttar", "extract a tar.gz to a buildlet", putTar)
	registerCommand("rm", "delete files or directories", rm)
	registerCommand("run", "run a command on a buildlet", run)
}

func main() {
	buildlet.RegisterFlags()
	registerCommands()
	flag.Usage = usage
	flag.Parse()
	buildEnv = buildenv.FromFlags()
	args := flag.Args()
	if len(args) == 0 {
		usage()
	}
	cmdName := args[0]
	cmd, ok := commands[cmdName]
	if !ok {
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmdName)
		usage()
	}
	err := cmd.run(args[1:])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error running %s: %v\n", cmdName, err)
		os.Exit(1)
	}
}
<commit_msg>cmd/gomote: update gomote example usage<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

/*
The gomote command is a client for the Go builder infrastructure.
It's a remote control for remote Go builder machines.

Usage:

  gomote [global-flags] cmd [cmd-flags]

  For example,
  $ gomote create openbsd-amd64-60
  user-username-openbsd-amd64-60-0
  $ gomote push user-username-openbsd-amd64-60-0
  $ gomote run user-username-openbsd-amd64-60-0 go/src/make.bash

TODO: document more, and figure out the CLI interface more.
*/
package main

import (
	"flag"
	"fmt"
	"os"
	"sort"

	"golang.org/x/build/buildenv"
	"golang.org/x/build/buildlet"
)

var (
	buildEnv *buildenv.Environment
)

type command struct {
	name string
	des  string
	run  func([]string) error
}

var commands = map[string]command{}

func sortedCommands() []string {
	s := make([]string, 0, len(commands))
	for name := range commands {
		s = append(s, name)
	}
	sort.Strings(s)
	return s
}

func usage() {
	fmt.Fprintf(os.Stderr, `Usage of gomote: gomote [global-flags] <cmd> [cmd-flags]

Global flags:
`)
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "Commands:\n\n")
	for _, name := range sortedCommands() {
		fmt.Fprintf(os.Stderr, "  %-10s %s\n", name, commands[name].des)
	}
	os.Exit(1)
}

func registerCommand(name, des string, run func([]string) error) {
	if _, dup := commands[name]; dup {
		panic("duplicate registration of " + name)
	}
	commands[name] = command{
		name: name,
		des:  des,
		run:  run,
	}
}

func registerCommands() {
	registerCommand("create", "create a buildlet", create)
	registerCommand("destroy", "destroy a buildlet", destroy)
	registerCommand("gettar", "extract a tar.gz from a buildlet", getTar)
	registerCommand("ls", "list the contents of a directory on a buildlet", ls)
	registerCommand("list", "list buildlets", list)
	registerCommand("ping", "test whether a buildlet is alive and reachable ", ping)
	registerCommand("push", "sync the repo of your pwd to the buildlet", push)
	registerCommand("put", "put files on a buildlet", put)
	registerCommand("put14", "put Go 1.4 in place", put14)
	registerCommand("puttar", "extract a tar.gz to a buildlet", putTar)
	registerCommand("rm", "delete files or directories", rm)
	registerCommand("run", "run a command on a buildlet", run)
}

func main() {
	buildlet.RegisterFlags()
	registerCommands()
	flag.Usage = usage
	flag.Parse()
	buildEnv = buildenv.FromFlags()
	args := flag.Args()
	if len(args) == 0 {
		usage()
	}
	cmdName := args[0]
	cmd, ok := commands[cmdName]
	if !ok {
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmdName)
		usage()
	}
	err := cmd.run(args[1:])
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error running %s: %v\n", cmdName, err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/juju/cmd"
	"launchpad.net/juju-core/juju/container"
	"launchpad.net/juju-core/juju/log"
	"launchpad.net/juju-core/juju/state"
	"launchpad.net/tomb"
)

// simpleContainer allows tests to hook into the container
// deployment logic.
var simpleContainer = container.Simple

// MachineAgent is a cmd.Command responsible for running a machine agent.
type MachineAgent struct {
	Conf      AgentConf
	MachineId int
}

// Info returns usage information for the command.
func (a *MachineAgent) Info() *cmd.Info {
	return &cmd.Info{"machine", "", "run a juju machine agent", ""}
}

// Init initializes the command for running.
func (a *MachineAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	f.IntVar(&a.MachineId, "machine-id", -1, "id of the machine to run")
	if err := f.Parse(true, args); err != nil {
		return err
	}
	if a.MachineId < 0 {
		return fmt.Errorf("--machine-id option must be set, and expects a non-negative integer")
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a machine agent.
func (a *MachineAgent) Run(_ *cmd.Context) error {
	// TODO reconnect when the agent fails.
	m, err := NewMachiner(&a.Conf.StateInfo, a.MachineId)
	if err != nil {
		return err
	}
	go m.loop()
	return m.Wait()
}

func NewMachiner(info *state.Info, machineId int) (*Machiner, error) {
	m := new(Machiner)
	var err error
	m.st, err = state.Open(info)
	if err != nil {
		return nil, err
	}
	m.machine, err = m.st.Machine(machineId)
	if err != nil {
		return nil, err
	}
	go m.loop()
	return m, nil
}

type Machiner struct {
	tomb    tomb.Tomb
	st      *state.State
	machine *state.Machine
}

func (m *Machiner) loop() {
	defer m.tomb.Done()
	defer m.st.Close()

	watcher := m.machine.WatchUnits()
	// TODO read initial units, check if they're running
	// and restart them if not.
	for {
		select {
		case <-m.tomb.Dying():
			return
		case change, ok := <-watcher.Changes():
			if !ok {
				err := watcher.Stop()
				if err != nil {
					m.tomb.Kill(err)
				}
				return
			}
			for _, u := range change.Deleted {
				if u.IsPrincipal() {
					if err := simpleContainer.Destroy(u); err != nil {
						// TODO what should we do in this case?
						log.Printf("cannot destroy unit %s: %v", u.Name(), err)
					}
				}
			}
			for _, u := range change.Added {
				if u.IsPrincipal() {
					if err := simpleContainer.Deploy(u); err != nil {
						// TODO what should we do in this case?
						log.Printf("cannot deploy unit %s: %v", u.Name(), err)
					}
				}
			}
		}
	}
}

func (m *Machiner) Wait() error {
	return m.tomb.Wait()
}

func (m *Machiner) Stop() error {
	m.tomb.Kill(nil)
	return m.tomb.Wait()
}
<commit_msg>cmd/jujud: cleaner stopping of machine agent<commit_after>package main

import (
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/juju/cmd"
	"launchpad.net/juju-core/juju/container"
	"launchpad.net/juju-core/juju/log"
	"launchpad.net/juju-core/juju/state"
	"launchpad.net/tomb"
)

// simpleContainer allows tests to hook into the container
// deployment logic.
var simpleContainer = container.Simple

// MachineAgent is a cmd.Command responsible for running a machine agent.
type MachineAgent struct {
	Conf      AgentConf
	MachineId int
}

// Info returns usage information for the command.
func (a *MachineAgent) Info() *cmd.Info {
	return &cmd.Info{"machine", "", "run a juju machine agent", ""}
}

// Init initializes the command for running.
func (a *MachineAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	f.IntVar(&a.MachineId, "machine-id", -1, "id of the machine to run")
	if err := f.Parse(true, args); err != nil {
		return err
	}
	if a.MachineId < 0 {
		return fmt.Errorf("--machine-id option must be set, and expects a non-negative integer")
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a machine agent.
func (a *MachineAgent) Run(_ *cmd.Context) error {
	// TODO reconnect when the agent fails.
	m, err := NewMachiner(&a.Conf.StateInfo, a.MachineId)
	if err != nil {
		return err
	}
	go m.loop()
	return m.Wait()
}

func NewMachiner(info *state.Info, machineId int) (*Machiner, error) {
	m := new(Machiner)
	var err error
	m.st, err = state.Open(info)
	if err != nil {
		return nil, err
	}
	m.machine, err = m.st.Machine(machineId)
	if err != nil {
		return nil, err
	}
	go m.loop()
	return m, nil
}

type Machiner struct {
	tomb    tomb.Tomb
	st      *state.State
	machine *state.Machine
}

func (m *Machiner) loop() {
	defer m.tomb.Done()
	defer m.st.Close()

	watcher := m.machine.WatchUnits()
	defer watcher.Stop()

	// TODO read initial units, check if they're running
	// and restart them if not.
	for {
		select {
		case <-m.tomb.Dying():
			return
		case change, ok := <-watcher.Changes():
			if !ok {
				err := watcher.Stop()
				if err == nil {
					panic("watcher closed channel with no error")
				}
				m.tomb.Kill(err)
				return
			}
			for _, u := range change.Deleted {
				if u.IsPrincipal() {
					if err := simpleContainer.Destroy(u); err != nil {
						// TODO what should we do in this case?
						log.Printf("cannot destroy unit %s: %v", u.Name(), err)
					}
				}
			}
			for _, u := range change.Added {
				if u.IsPrincipal() {
					if err := simpleContainer.Deploy(u); err != nil {
						// TODO what should we do in this case?
						log.Printf("cannot deploy unit %s: %v", u.Name(), err)
					}
				}
			}
		}
	}
}

func (m *Machiner) Wait() error {
	return m.tomb.Wait()
}

func (m *Machiner) Stop() error {
	m.tomb.Kill(nil)
	return m.tomb.Wait()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"http"
	"junta/client"
	"junta/mon"
	"junta/paxos"
	"junta/server"
	"junta/store"
	"junta/util"
	"junta/web"
	"net"
	"os"
	"strings"
)

const (
	alpha = 50
)

// Flags
var (
	listenAddr = flag.String("l", "127.0.0.1:8046", "The address to bind to.")
	publishAddr = flag.String("L", "", "The address puslished for remote clients (default is listen address)")
	attachAddr = flag.String("a", "", "The address of another node to attach to.")
	webAddr = flag.String("w", ":8080", "Serve web requests on this address.")
	clusterName = flag.String("c", "local", "The non-empty cluster name.")
)

func activate(st *store.Store, self, prefix string, c *client.Client) {
	logger := util.NewLogger("activate")
	ch := make(chan store.Event)
	st.Watch("/junta/slot/*", ch)

	go func() {
		for _, slot := range store.GetDir(st, "/junta/slot") {
			p := "/junta/slot/" + slot
			v, cas := st.Get(p)
			ch <- store.Event{0, p, v[0], cas, "", nil, nil}
		}
	}()

	for ev := range ch {
		// TODO ev.IsEmpty()
		if ev.IsSet() && ev.Body == "" {
			_, err := c.Set(prefix+ev.Path, self, ev.Cas)
			if err == nil {
				return
			}
			logger.Log(err)
		}
	}
}

func Usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS] <cluster-name>\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\nOptions:\n")
	flag.PrintDefaults()
}

func main() {
	util.LogWriter = os.Stderr
	logger := util.NewLogger("main")

	flag.Parse()
	flag.Usage = Usage

	prefix := "/j/" + *clusterName

	if *listenAddr == "" {
		logger.Log("require a listen address")
		flag.Usage()
		os.Exit(1)
	}

	if *publishAddr == "" {
		publishAddr = listenAddr
	}

	var webListener net.Listener
	if *webAddr != "" {
		wl, err := net.Listen("tcp", *webAddr)
		if err != nil {
			panic(err)
		}
		webListener = wl
	}

	listener, err := net.Listen("tcp", *listenAddr)
	if err != nil {
		panic(err)
	}

	publishParts := strings.Split(*publishAddr, ":", 2)
	if len(publishParts) < 2 && publishParts[0] == "" {
		logger.Log("invalid publish address")
		flag.Usage()
		os.Exit(1)
	}

	outs := make(paxos.ChanPutCloserTo)

	var cl *client.Client
	self := util.RandId()
	st := store.New()
	seqn := uint64(0)
	if *attachAddr == "" { // we are the only node in a new cluster
		seqn = addPublicAddr(st, seqn + 1, self, *listenAddr)
		seqn = addHostname(st, seqn + 1, self, publishParts[0])
		seqn = addMember(st, seqn + 1, self, *listenAddr)
		seqn = claimSlot(st, seqn + 1, "1", self)
		seqn = claimLeader(st, seqn + 1, self)
		seqn = claimSlot(st, seqn + 1, "2", "")
		seqn = claimSlot(st, seqn + 1, "3", "")
		seqn = claimSlot(st, seqn + 1, "4", "")
		seqn = claimSlot(st, seqn + 1, "5", "")
		seqn = addPing(st, seqn + 1, "pong")

		cl, err = client.Dial(*listenAddr)
		if err != nil {
			panic(err)
		}
	} else {
		cl, err = client.Dial(*attachAddr)
		if err != nil {
			panic(err)
		}

		path := prefix + "/junta/info/"+ self +"/public-addr"
		_, err = cl.Set(path, *listenAddr, store.Clobber)
		if err != nil {
			panic(err)
		}

		path = prefix + "/junta/info/"+ self +"/hostname"
		_, err = cl.Set(path, os.Getenv("HOSTNAME"), store.Clobber)
		if err != nil {
			panic(err)
		}

		var snap string
		seqn, snap, err = cl.Join(self, *listenAddr)
		if err != nil {
			panic(err)
		}

		ch := make(chan store.Event)
		st.Wait(seqn + alpha, ch)
		st.Apply(1, snap)

		go func() {
			<-ch
			activate(st, self, prefix, cl)
		}()

		// TODO sink needs a way to pick up missing values if there are any
		// gaps in its sequence
	}
	mg := paxos.NewManager(self, seqn, alpha, st, outs)

	if *attachAddr == "" {
		// Skip ahead alpha steps so that the registrar can provide a
		// meaningful cluster.
		for i := seqn + 1; i < seqn + alpha; i++ {
			go st.Apply(i, store.Nop)
		}
	}

	sv := &server.Server{*listenAddr, st, mg, self, prefix}

	go func() {
		panic(mon.Monitor(self, prefix, st, cl))
	}()

	go func() {
		panic(sv.Serve(listener))
	}()

	go func() {
		panic(sv.ListenAndServeUdp(outs))
	}()

	if webListener != nil {
		web.Store = st
		web.MainInfo.ClusterName = *clusterName
		// http handlers are installed in the init function of junta/web.
		go http.Serve(webListener, nil)
	}

	for {
		st.Apply(mg.Recv())
	}
}

func addPublicAddr(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/junta/info/"+ self +"/public-addr"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func addHostname(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/junta/info/"+ self +"/hostname"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func addMember(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/junta/members/"+self, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func claimSlot(st *store.Store, seqn uint64, slot, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/junta/slot/"+slot, self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func claimLeader(st *store.Store, seqn uint64, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/junta/leader", self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func addPing(st *store.Store, seqn uint64, v string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/ping", v, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}
<commit_msg>use sugar<commit_after>package main

import (
	"flag"
	"fmt"
	"http"
	"junta/client"
	"junta/mon"
	"junta/paxos"
	"junta/server"
	"junta/store"
	"junta/util"
	"junta/web"
	"net"
	"os"
	"strings"
)

const (
	alpha = 50
)

// Flags
var (
	listenAddr = flag.String("l", "127.0.0.1:8046", "The address to bind to.")
	publishAddr = flag.String("L", "", "The address puslished for remote clients (default is listen address)")
	attachAddr = flag.String("a", "", "The address of another node to attach to.")
	webAddr = flag.String("w", ":8080", "Serve web requests on this address.")
	clusterName = flag.String("c", "local", "The non-empty cluster name.")
)

func activate(st *store.Store, self, prefix string, c *client.Client) {
	logger := util.NewLogger("activate")
	ch := make(chan store.Event)
	st.GetDirAndWatch("/junta/slot", ch)
	for ev := range ch {
		// TODO ev.IsEmpty()
		if ev.IsSet() && ev.Body == "" {
			_, err := c.Set(prefix+ev.Path, self, ev.Cas)
			if err == nil {
				return
			}
			logger.Log(err)
		}
	}
}

func Usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS] <cluster-name>\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\nOptions:\n")
	flag.PrintDefaults()
}

func main() {
	util.LogWriter = os.Stderr
	logger := util.NewLogger("main")

	flag.Parse()
	flag.Usage = Usage

	prefix := "/j/" + *clusterName

	if *listenAddr == "" {
		logger.Log("require a listen address")
		flag.Usage()
		os.Exit(1)
	}

	if *publishAddr == "" {
		publishAddr = listenAddr
	}

	var webListener net.Listener
	if *webAddr != "" {
		wl, err := net.Listen("tcp", *webAddr)
		if err != nil {
			panic(err)
		}
		webListener = wl
	}

	listener, err := net.Listen("tcp", *listenAddr)
	if err != nil {
		panic(err)
	}

	publishParts := strings.Split(*publishAddr, ":", 2)
	if len(publishParts) < 2 && publishParts[0] == "" {
		logger.Log("invalid publish address")
		flag.Usage()
		os.Exit(1)
	}

	outs := make(paxos.ChanPutCloserTo)

	var cl *client.Client
	self := util.RandId()
	st := store.New()
	seqn := uint64(0)
	if *attachAddr == "" { // we are the only node in a new cluster
		seqn = addPublicAddr(st, seqn + 1, self, *listenAddr)
		seqn = addHostname(st, seqn + 1, self, publishParts[0])
		seqn = addMember(st, seqn + 1, self, *listenAddr)
		seqn = claimSlot(st, seqn + 1, "1", self)
		seqn = claimLeader(st, seqn + 1, self)
		seqn = claimSlot(st, seqn + 1, "2", "")
		seqn = claimSlot(st, seqn + 1, "3", "")
		seqn = claimSlot(st, seqn + 1, "4", "")
		seqn = claimSlot(st, seqn + 1, "5", "")
		seqn = addPing(st, seqn + 1, "pong")

		cl, err = client.Dial(*listenAddr)
		if err != nil {
			panic(err)
		}
	} else {
		cl, err = client.Dial(*attachAddr)
		if err != nil {
			panic(err)
		}

		path := prefix + "/junta/info/"+ self +"/public-addr"
		_, err = cl.Set(path, *listenAddr, store.Clobber)
		if err != nil {
			panic(err)
		}

		path = prefix + "/junta/info/"+ self +"/hostname"
		_, err = cl.Set(path, os.Getenv("HOSTNAME"), store.Clobber)
		if err != nil {
			panic(err)
		}

		var snap string
		seqn, snap, err = cl.Join(self, *listenAddr)
		if err != nil {
			panic(err)
		}

		ch := make(chan store.Event)
		st.Wait(seqn + alpha, ch)
		st.Apply(1, snap)

		go func() {
			<-ch
			activate(st, self, prefix, cl)
		}()

		// TODO sink needs a way to pick up missing values if there are any
		// gaps in its sequence
	}
	mg := paxos.NewManager(self, seqn, alpha, st, outs)

	if *attachAddr == "" {
		// Skip ahead alpha steps so that the registrar can provide a
		// meaningful cluster.
		for i := seqn + 1; i < seqn + alpha; i++ {
			go st.Apply(i, store.Nop)
		}
	}

	sv := &server.Server{*listenAddr, st, mg, self, prefix}

	go func() {
		panic(mon.Monitor(self, prefix, st, cl))
	}()

	go func() {
		panic(sv.Serve(listener))
	}()

	go func() {
		panic(sv.ListenAndServeUdp(outs))
	}()

	if webListener != nil {
		web.Store = st
		web.MainInfo.ClusterName = *clusterName
		// http handlers are installed in the init function of junta/web.
		go http.Serve(webListener, nil)
	}

	for {
		st.Apply(mg.Recv())
	}
}

func addPublicAddr(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/junta/info/"+ self +"/public-addr"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func addHostname(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/junta/info/"+ self +"/hostname"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func addMember(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/junta/members/"+self, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func claimSlot(st *store.Store, seqn uint64, slot, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/junta/slot/"+slot, self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func claimLeader(st *store.Store, seqn uint64, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/junta/leader", self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}

func addPing(st *store.Store, seqn uint64, v string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/ping", v, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"sort"
	"strconv"

	_ "code.google.com/p/go-sqlite/go1/sqlite3"
	"github.com/rwcarlsen/cyan/query"
)

var (
	help   = flag.Bool("h", false, "Print this help message")
	dbname = flag.String("db", "", "cyclus sqlite database to query")
	simid  = flag.String("simid", "", "simulation id (empty string defaults to first sim id in database")
)

var command string

var db *sql.DB

var cmds = NewCmdSet()

func init() {
	cmds.Register("agents", "list all agents in the simulation", doAgents)
	cmds.Register("sims", "list all simulations in the database", doSims)
	cmds.Register("inv", "show inventory of one or more agents at a specific timestep", doInv)
	cmds.Register("created", "show material created by one or more agents between specific timesteps", doCreated)
	cmds.Register("deployseries", "print a time-series of a prototype's total active deployments", doDeploySeries)
	cmds.Register("flow", "Show total transacted material between two groups of agents between specific timesteps", doFlow)
	cmds.Register("invseries", "print a time series of an agent's inventory for specified isotopes", doInvSeries)
	cmds.Register("flowgraph", "print a graphviz dot graph of resource arcs between facilities", doFlowGraph)
}

func main() {
	log.SetFlags(0)
	flag.Parse()

	if *help || flag.NArg() < 1 {
		fmt.Println("Usage: metric -db <cyclus-db> [opts] <command> [args...]")
		fmt.Println("Calculates metrics for cyclus simulation data in a sqlite database.")
		flag.PrintDefaults()
		fmt.Println("\nCommands:")
		for i := range cmds.Names {
			fmt.Printf("    - %v: %v\n", cmds.Names[i], cmds.Helps[i])
		}
		return
	} else if *dbname == "" {
		log.Fatal("must specify database with db flag")
	}

	var err error
	db, err = sql.Open("sqlite3", *dbname)
	fatalif(err)
	defer db.Close()

	if *simid == "" {
		ids, err := query.SimIds(db)
		fatalif(err)
		*simid = ids[0]
	}

	cmds.Execute(flag.Args())
}

func doSims(args []string) {
	ids, err := query.SimIds(db)
	fatalif(err)
	for _, id := range ids {
		info, err := query.SimStat(db, id)
		fatalif(err)
		fmt.Println(info)
	}
}

func doAgents(args []string) {
	ags, err := query.AllAgents(db, *simid)
	fatalif(err)
	for _, a := range ags {
		fmt.Println(a)
	}
}

func doInv(args []string) {
	fs := flag.NewFlagSet("inv", flag.ExitOnError)
	t := fs.Int("t", -1, "timestep of inventory (-1 = end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: inv [agent-id...]\nZero agents uses all agent inventories")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	var agents []int
	for _, arg := range fs.Args() {
		id, err := strconv.Atoi(arg)
		fatalif(err)
		agents = append(agents, id)
	}

	m, err := query.InvAt(db, *simid, *t, agents...)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

type Row struct {
	X  int
	Ys []float64
}

type MultiSeries [][]query.XY

func (ms MultiSeries) Rows() []Row {
	rowmap := map[int]Row{}
	xs := []int{}
	for i, s := range ms {
		for _, xy := range s {
			row, ok := rowmap[xy.X]
			if !ok {
				xs = append(xs, xy.X)
				row.Ys = make([]float64, len(ms))
				row.X = xy.X
			}
			row.Ys[i] = xy.Y
			rowmap[xy.X] = row
		}
	}

	sort.Ints(xs)
	rows := make([]Row, 0, len(rowmap))
	for _, x := range xs {
		rows = append(rows, rowmap[x])
	}
	return rows
}

func doInvSeries(args []string) {
	fs := flag.NewFlagSet("invseries", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: invseries <agent-id> <isotope> [isotope...]"); fs.PrintDefaults() }
	fs.Parse(args)
	if fs.NArg() < 2 {
		fs.Usage()
		return
	}

	agent, err := strconv.Atoi(fs.Arg(0))
	fatalif(err)

	isos := []int{}
	for _, arg := range fs.Args()[1:] {
		iso, err := strconv.Atoi(arg)
		fatalif(err)
		isos = append(isos, iso)
	}

	ms := MultiSeries{}
	for _, iso := range isos {
		xys, err := query.InvSeries(db, *simid, agent, iso)
		ms = append(ms, xys)
		fatalif(err)
	}

	fmt.Printf("# Agent %v inventory in kg\n", agent)
	fmt.Printf("# [Timestep]")
	for _, iso := range isos {
		fmt.Printf(" [%v]", iso)
	}
	fmt.Printf("\n")
	for _, row := range ms.Rows() {
		fmt.Printf("%v", row.X)
		for _, y := range row.Ys {
			fmt.Printf(" %v ", y)
		}
		fmt.Printf("\n")
	}
}

func doFlowGraph(args []string) {
	fs := flag.NewFlagSet("flowgraph", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: flowgraph"); fs.PrintDefaults() }
	proto := fs.Bool("proto", false, "aggregate nodes by prototype")
	t0 := fs.Int("t1", -1, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Parse(args)

	arcs, err := query.FlowGraph(db, *simid, *t0, *t1, *proto)
	fatalif(err)

	fmt.Println("digraph ResourceFlows {")
	for _, arc := range arcs {
		fmt.Printf("    \"%v\" -> \"%v\" [label=\"%v (%v kg)\", fontsize=8];\n", arc.Src, arc.Dst, arc.Commod, arc.Quantity)
	}
	fmt.Println("}")
}

func doDeploySeries(args []string) {
	fs := flag.NewFlagSet("deployseries", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: deployseries <prototype>"); fs.PrintDefaults() }
	fs.Parse(args)
	if fs.NArg() < 1 {
		fs.Usage()
		return
	}

	proto := fs.Arg(0)
	xys, err := query.DeployCumulative(db, *simid, proto)
	fatalif(err)

	fmt.Printf("# Prototype %v total active deployments\n", proto)
	fmt.Println("# [Timestep] [Count]")
	for _, xy := range xys {
		fmt.Printf("%v %v\n", xy.X, xy.Y)
	}
}

func doCreated(args []string) {
	fs := flag.NewFlagSet("created", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: created [agent-id...]\nZero agents uses all agents"); fs.PrintDefaults() }
	t0 := fs.Int("t1", -1, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Parse(args)

	var agents []int

	for _, arg := range fs.Args() {
		id, err := strconv.Atoi(arg)
		fatalif(err)
		agents = append(agents, id)
	}

	m, err := query.MatCreated(db, *simid, *t0, *t1, agents...)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

func doFlow(args []string) {
	fs := flag.NewFlagSet("created", flag.ExitOnError)
	t0 := fs.Int("t1", -1, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: flow <from-agents...> .. <to-agents...>\nZero agents uses all agents")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	var from []int
	var to []int

	if flag.NArg() < 3 {
		fs.Usage()
		return
	}

	before := true
	for _, arg := range flag.Args()[3:] {
		if arg == ".." {
			before = false
			continue
		}

		id, err := strconv.Atoi(arg)
		fatalif(err)
		if before {
			from = append(from, id)
		} else {
			to = append(to, id)
		}
	}
	if len(to) < 1 {
		fs.Usage()
		return
	}

	m, err := query.Flow(db, *simid, *t0, *t1, from, to)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

func fatalif(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

type CmdSet struct {
	funcs map[string]func([]string)
	Names []string
	Helps []string
}

func NewCmdSet() *CmdSet {
	return &CmdSet{funcs: map[string]func([]string){}}
}

func (cs *CmdSet) Register(name, brief string, f func([]string)) {
	cs.Names = append(cs.Names, name)
	cs.Helps = append(cs.Helps, brief)
	cs.funcs[name] = f
}

func (cs *CmdSet) Execute(args []string) {
	cmd := args[0]
	f, ok := cs.funcs[cmd]
	if !ok {
		log.Fatalf("Invalid command '%v'", cmd)
	}
	f(args[1:])
}
<commit_msg>tweaked format<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"sort"
	"strconv"

	_ "code.google.com/p/go-sqlite/go1/sqlite3"
	"github.com/rwcarlsen/cyan/query"
)

var (
	help   = flag.Bool("h", false, "Print this help message")
	dbname = flag.String("db", "", "cyclus sqlite database to query")
	simid  = flag.String("simid", "", "simulation id (empty string defaults to first sim id in database")
)

var command string

var db *sql.DB

var cmds = NewCmdSet()

func init() {
	cmds.Register("agents", "list all agents in the simulation", doAgents)
	cmds.Register("sims", "list all simulations in the database", doSims)
	cmds.Register("inv", "show inventory of one or more agents at a specific timestep", doInv)
	cmds.Register("created", "show material created by one or more agents between specific timesteps", doCreated)
	cmds.Register("deployseries", "print a time-series of a prototype's total active deployments", doDeploySeries)
	cmds.Register("flow", "Show total transacted material between two groups of agents between specific timesteps", doFlow)
	cmds.Register("invseries", "print a time series of an agent's inventory for specified isotopes", doInvSeries)
	cmds.Register("flowgraph", "print a graphviz dot graph of resource arcs between facilities", doFlowGraph)
}

func main() {
	log.SetFlags(0)
	flag.Parse()

	if *help || flag.NArg() < 1 {
		fmt.Println("Usage: metric -db <cyclus-db> [opts] <command> [args...]")
		fmt.Println("Calculates metrics for cyclus simulation data in a sqlite database.")
		flag.PrintDefaults()
		fmt.Println("\nCommands:")
		for i := range cmds.Names {
			fmt.Printf("    - %v: %v\n", cmds.Names[i], cmds.Helps[i])
		}
		return
	} else if *dbname == "" {
		log.Fatal("must specify database with db flag")
	}

	var err error
	db, err = sql.Open("sqlite3", *dbname)
	fatalif(err)
	defer db.Close()

	if *simid == "" {
		ids, err := query.SimIds(db)
		fatalif(err)
		*simid = ids[0]
	}

	cmds.Execute(flag.Args())
}

func doSims(args []string) {
	ids, err := query.SimIds(db)
	fatalif(err)
	for _, id := range ids {
		info, err := query.SimStat(db, id)
		fatalif(err)
		fmt.Println(info)
	}
}

func doAgents(args []string) {
	ags, err := query.AllAgents(db, *simid)
	fatalif(err)
	for _, a := range ags {
		fmt.Println(a)
	}
}

func doInv(args []string) {
	fs := flag.NewFlagSet("inv", flag.ExitOnError)
	t := fs.Int("t", -1, "timestep of inventory (-1 = end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: inv [agent-id...]\nZero agents uses all agent inventories")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	var agents []int
	for _, arg := range fs.Args() {
		id, err := strconv.Atoi(arg)
		fatalif(err)
		agents = append(agents, id)
	}

	m, err := query.InvAt(db, *simid, *t, agents...)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

type Row struct {
	X  int
	Ys []float64
}

type MultiSeries [][]query.XY

func (ms MultiSeries) Rows() []Row {
	rowmap := map[int]Row{}
	xs := []int{}
	for i, s := range ms {
		for _, xy := range s {
			row, ok := rowmap[xy.X]
			if !ok {
				xs = append(xs, xy.X)
				row.Ys = make([]float64, len(ms))
				row.X = xy.X
			}
			row.Ys[i] = xy.Y
			rowmap[xy.X] = row
		}
	}

	sort.Ints(xs)
	rows := make([]Row, 0, len(rowmap))
	for _, x := range xs {
		rows = append(rows, rowmap[x])
	}
	return rows
}

func doInvSeries(args []string) {
	fs := flag.NewFlagSet("invseries", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: invseries <agent-id> <isotope> [isotope...]"); fs.PrintDefaults() }
	fs.Parse(args)
	if fs.NArg() < 2 {
		fs.Usage()
		return
	}

	agent, err := strconv.Atoi(fs.Arg(0))
	fatalif(err)

	isos := []int{}
	for _, arg := range fs.Args()[1:] {
		iso, err := strconv.Atoi(arg)
		fatalif(err)
		isos = append(isos, iso)
	}

	ms := MultiSeries{}
	for _, iso := range isos {
		xys, err := query.InvSeries(db, *simid, agent, iso)
		ms = append(ms, xys)
		fatalif(err)
	}

	fmt.Printf("# Agent %v inventory in kg\n", agent)
	fmt.Printf("# [Timestep]")
	for _, iso := range isos {
		fmt.Printf(" [%v]", iso)
	}
	fmt.Printf("\n")
	for _, row := range ms.Rows() {
		fmt.Printf("%v", row.X)
		for _, y := range row.Ys {
			fmt.Printf(" %v ", y)
		}
		fmt.Printf("\n")
	}
}

func doFlowGraph(args []string) {
	fs := flag.NewFlagSet("flowgraph", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: flowgraph"); fs.PrintDefaults() }
	proto := fs.Bool("proto", false, "aggregate nodes by prototype")
	t0 := fs.Int("t1", -1, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Parse(args)

	arcs, err := query.FlowGraph(db, *simid, *t0, *t1, *proto)
	fatalif(err)

	fmt.Println("digraph ResourceFlows {")
	fmt.Println("    overlap = false;")
	fmt.Println("    nodesep=1.0;")
	fmt.Println("    edge [fontsize=9];")
	for _, arc := range arcs {
		fmt.Printf("    \"%v\" -> \"%v\" [label=\"%v\n(%.3g kg)\"];\n", arc.Src, arc.Dst, arc.Commod, arc.Quantity)
	}
	fmt.Println("}")
}

func doDeploySeries(args []string) {
	fs := flag.NewFlagSet("deployseries", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: deployseries <prototype>"); fs.PrintDefaults() }
	fs.Parse(args)
	if fs.NArg() < 1 {
		fs.Usage()
		return
	}

	proto := fs.Arg(0)
	xys, err := query.DeployCumulative(db, *simid, proto)
	fatalif(err)

	fmt.Printf("# Prototype %v total active deployments\n", proto)
	fmt.Println("# [Timestep] [Count]")
	for _, xy := range xys {
		fmt.Printf("%v %v\n", xy.X, xy.Y)
	}
}

func doCreated(args []string) {
	fs := flag.NewFlagSet("created", flag.ExitOnError)
	fs.Usage = func() { log.Print("Usage: created [agent-id...]\nZero agents uses all agents"); fs.PrintDefaults() }
	t0 := fs.Int("t1", -1, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Parse(args)

	var agents []int

	for _, arg := range fs.Args() {
		id, err := strconv.Atoi(arg)
		fatalif(err)
		agents = append(agents, id)
	}

	m, err := query.MatCreated(db, *simid, *t0, *t1, agents...)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

func doFlow(args []string) {
	fs := flag.NewFlagSet("created", flag.ExitOnError)
	t0 := fs.Int("t1", -1, "beginning of time interval (default is beginning of simulation)")
	t1 := fs.Int("t2", -1, "end of time interval (default if end of simulation)")
	fs.Usage = func() {
		log.Print("Usage: flow <from-agents...> .. <to-agents...>\nZero agents uses all agents")
		fs.PrintDefaults()
	}
	fs.Parse(args)

	var from []int
	var to []int

	if flag.NArg() < 3 {
		fs.Usage()
		return
	}

	before := true
	for _, arg := range flag.Args()[3:] {
		if arg == ".." {
			before = false
			continue
		}

		id, err := strconv.Atoi(arg)
		fatalif(err)
		if before {
			from = append(from, id)
		} else {
			to = append(to, id)
		}
	}
	if len(to) < 1 {
		fs.Usage()
		return
	}

	m, err := query.Flow(db, *simid, *t0, *t1, from, to)
	fatalif(err)
	fmt.Printf("%+v\n", m)
}

func fatalif(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

type CmdSet struct {
	funcs map[string]func([]string)
	Names []string
	Helps []string
}

func NewCmdSet() *CmdSet {
	return &CmdSet{funcs: map[string]func([]string){}}
}

func (cs *CmdSet) Register(name, brief string, f func([]string)) {
	cs.Names = append(cs.Names, name)
	cs.Helps = append(cs.Helps, brief)
	cs.funcs[name] = f
}

func (cs *CmdSet) Execute(args []string) {
	cmd := args[0]
	f, ok := cs.funcs[cmd]
	if !ok {
		log.Fatalf("Invalid command '%v'", cmd)
	}
	f(args[1:])
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"

	"github.com/cbroglie/mustache"
)

var rootCmd = &cobra.Command{
	Use: "mustache [--layout template] [data] template",
	Example: `  $ mustache data.yml template.mustache
  $ cat data.yml | mustache template.mustache
  $ mustache --layout wrapper.mustache data template.mustache
  $ mustache --overide over.yml data.yml template.mustache`,
	Args: cobra.RangeArgs(0, 2),
	Run: func(cmd *cobra.Command, args []string) {
		err := run(cmd, args)
		if err != nil {
			fmt.Printf("Error: %s\n", err.Error())
			os.Exit(1)
		}
	},
}
var layoutFile string
var overrideFile string

func main() {
	rootCmd.Flags().StringVar(&layoutFile, "layout", "", "location of layout file")
	rootCmd.Flags().StringVar(&overrideFile, "override", "", "location of data.yml override yml")
	if err := rootCmd.Execute(); err != nil {
		os.Exit(1)
	}
}

func run(cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		cmd.Usage()
		return nil
	}

	var data interface{}
	var templatePath string
	if len(args) == 1 {
		var err error
		data, err = parseDataFromStdIn()
		if err != nil {
			return err
		}
		templatePath = args[0]
	} else {
		var err error
		data, err = parseDataFromFile(args[0])
		if err != nil {
			return err
		}
		templatePath = args[1]
	}

	if overrideFile != "" {
		override, err := parseDataFromFile(overrideFile)
		if err != nil {
			return err
		}
		for k, v := range override.(map[interface{}]interface{}) {
			data.(map[interface{}]interface{})[k] = v
		}
	}
	var output string
	var err error
	if layoutFile != "" {
		output, err = mustache.RenderFileInLayout(templatePath, layoutFile, data)
	} else {
		output, err = mustache.RenderFile(templatePath, data)
	}
	if err != nil {
		return err
	}
	fmt.Print(output)
	return nil
}

func parseDataFromStdIn() (interface{}, error) {
	b, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := yaml.Unmarshal(b, &data); err != nil {
		return nil, err
	}
	return data, nil
}

func parseDataFromFile(filePath string) (interface{}, error) {
	b, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := yaml.Unmarshal(b, &data); err != nil {
		return nil, err
	}
	return data, nil
}
<commit_msg>Fix typo in help text<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"

	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"

	"github.com/cbroglie/mustache"
)

var rootCmd = &cobra.Command{
	Use: "mustache [--layout template] [data] template",
	Example: `  $ mustache data.yml template.mustache
  $ cat data.yml | mustache template.mustache
  $ mustache --layout wrapper.mustache data template.mustache
  $ mustache --override over.yml data.yml template.mustache`,
	Args: cobra.RangeArgs(0, 2),
	Run: func(cmd *cobra.Command, args []string) {
		err := run(cmd, args)
		if err != nil {
			fmt.Printf("Error: %s\n", err.Error())
			os.Exit(1)
		}
	},
}
var layoutFile string
var overrideFile string

func main() {
	rootCmd.Flags().StringVar(&layoutFile, "layout", "", "location of layout file")
	rootCmd.Flags().StringVar(&overrideFile, "override", "", "location of data.yml override yml")
	if err := rootCmd.Execute(); err != nil {
		os.Exit(1)
	}
}

func run(cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		cmd.Usage()
		return nil
	}

	var data interface{}
	var templatePath string
	if len(args) == 1 {
		var err error
		data, err = parseDataFromStdIn()
		if err != nil {
			return err
		}
		templatePath = args[0]
	} else {
		var err error
		data, err = parseDataFromFile(args[0])
		if err != nil {
			return err
		}
		templatePath = args[1]
	}

	if overrideFile != "" {
		override, err := parseDataFromFile(overrideFile)
		if err != nil {
			return err
		}
		for k, v := range override.(map[interface{}]interface{}) {
			data.(map[interface{}]interface{})[k] = v
		}
	}
	var output string
	var err error
	if layoutFile != "" {
		output, err = mustache.RenderFileInLayout(templatePath, layoutFile, data)
	} else {
		output, err = mustache.RenderFile(templatePath, data)
	}
	if err != nil {
		return err
	}
	fmt.Print(output)
	return nil
}

func parseDataFromStdIn() (interface{}, error) {
	b, err := ioutil.ReadAll(os.Stdin)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := yaml.Unmarshal(b, &data); err != nil {
		return nil, err
	}
	return data, nil
}

func parseDataFromFile(filePath string) (interface{}, error) {
	b, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := yaml.Unmarshal(b, &data); err != nil {
		return nil, err
	}
	return data, nil
}
<|endoftext|>