<commit_before>// Copyright © 2020 Karim Radhouani <medkarimrdi@gmail.com>
//
// 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 cmd

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	_ "github.com/karimra/gnmic/outputs/all"
	"github.com/manifoldco/promptui"
	"github.com/mitchellh/mapstructure"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const defaultRetryTimer = 10 * time.Second

var subscriptionModes = [][2]string{
	{"once", "a single request/response channel. The target creates the relevant update messages, transmits them, and subsequently closes the RPC"},
	{"stream", "long-lived subscriptions which continue to transmit updates relating to the set of paths that are covered within the subscription indefinitely"},
	{"poll", "on-demand retrieval of data items via long-lived RPCs"},
}
var streamSubscriptionModes = [][2]string{
	{"target-defined", "the target MUST determine the best type of subscription to be created on a per-leaf basis"},
	{"sample", "the value of the data item(s) MUST be sent once per sample interval to the client"},
	{"on-change", "data updates are only sent when the value of the data item changes"},
}

// subscribeCmd represents the subscribe command
var subscribeCmd = &cobra.Command{
	Use:     "subscribe",
	Aliases: []string{"sub"},
	Short:   "subscribe to gnmi updates on targets",
	Annotations: map[string]string{
		"--path":        "XPATH",
		"--prefix":      "PREFIX",
		"--model":       "MODEL",
		"--mode":        "SUBSC_MODE",
		"--stream-mode": "STREAM_MODE",
	},

	RunE: func(cmd *cobra.Command, args []string) error {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		setupCloseHandler(cancel)
		debug := viper.GetBool("debug")
		targetsConfig, err := createTargets()
		if err != nil {
			return fmt.Errorf("failed getting targets config: %v", err)
		}
		if debug {
			logger.Printf("targets: %s", targetsConfig)
		}
		subscriptionsConfig, err := getSubscriptions()
		if err != nil {
			return fmt.Errorf("failed getting subscriptions config: %v", err)
		}
		if debug {
			logger.Printf("subscriptions: %s", subscriptionsConfig)
		}
		outs, err := getOutputs(ctx)
		if err != nil {
			return err
		}
		if debug {
			logger.Printf("outputs: %+v", outs)
		}

		cfg := &collector.Config{
			PrometheusAddress:   viper.GetString("prometheus-address"),
			Debug:               viper.GetBool("debug"),
			Format:              viper.GetString("format"),
			TargetReceiveBuffer: viper.GetUint("target-buffer-size"),
			RetryTimer:          viper.GetDuration("retry-timer"),
		}

		coll := collector.NewCollector(cfg, targetsConfig, subscriptionsConfig, outs, createCollectorDialOpts(), logger)

		wg := new(sync.WaitGroup)
		wg.Add(len(coll.Targets))
		for name := range coll.Targets {
			go func(tn string) {
				defer wg.Done()
				tRetryTimer := coll.Targets[tn].Config.RetryTimer
				for {
					err = coll.Subscribe(ctx, tn)
					if err != nil {
						if errors.Is(err, context.DeadlineExceeded) {
							logger.Printf("failed to initialize target '%s' timeout (%s) reached", tn, targetsConfig[tn].Timeout)
						} else {
							logger.Printf("failed to initialize target '%s': %v", tn, err)
						}
						logger.Printf("retrying target %s in %s", tn, tRetryTimer)
						time.Sleep(tRetryTimer)
						continue
					}
					return
				}
			}(name)
		}
		wg.Wait()
		polledTargetsSubscriptions := coll.PolledSubscriptionsTargets()
		if len(polledTargetsSubscriptions) > 0 {
			pollTargets := make([]string, 0, len(polledTargetsSubscriptions))
			for t := range polledTargetsSubscriptions {
				pollTargets = append(pollTargets, t)
			}
			sort.Slice(pollTargets, func(i, j int) bool {
				return pollTargets[i] < pollTargets[j]
			})
			s := promptui.Select{
				Label:        "select target to poll",
				Items:        pollTargets,
				HideSelected: true,
			}
			waitChan := make(chan struct{}, 1)
			waitChan <- struct{}{}
			mo := &collector.MarshalOptions{
				Multiline: true,
				Indent:    "  ",
				Format:    viper.GetString("format")}
			go func() {
				for {
					select {
					case <-waitChan:
						_, name, err := s.Run()
						if err != nil {
							fmt.Printf("failed selecting target to poll: %v\n", err)
							continue
						}
						ss := promptui.Select{
							Label:        "select subscription to poll",
							Items:        polledTargetsSubscriptions[name],
							HideSelected: true,
						}
						_, subName, err := ss.Run()
						if err != nil {
							fmt.Printf("failed selecting subscription to poll: %v\n", err)
							continue
						}
						response, err := coll.TargetPoll(name, subName)
						if err != nil && err != io.EOF {
							fmt.Printf("target '%s', subscription '%s': poll response error:%v\n", name, subName, err)
							continue
						}
						if response == nil {
							fmt.Printf("received empty response from target '%s'\n", name)
							continue
						}
						switch rsp := response.Response.(type) {
						case *gnmi.SubscribeResponse_SyncResponse:
							fmt.Printf("received sync response '%t' from '%s'\n", rsp.SyncResponse, name)
							waitChan <- struct{}{}
							continue
						}
						b, err := mo.Marshal(response, nil)
						if err != nil {
							fmt.Printf("target '%s', subscription '%s': poll response formatting error:%v\n", name, subName, err)
							fmt.Println(string(b))
							waitChan <- struct{}{}
							continue
						}
						fmt.Println(string(b))
						waitChan <- struct{}{}
					case <-ctx.Done():
						return
					}
				}
			}()
		}
		coll.Start(ctx)
		return nil
	},
	PostRun: func(cmd *cobra.Command, args []string) {
		cmd.ResetFlags()
		initSubscribeFlags(cmd)
	},
	SilenceUsage: true,
}

func init() {
	rootCmd.AddCommand(subscribeCmd)
	initSubscribeFlags(subscribeCmd)
}

// used to init or reset subscribeCmd flags for gnmic-prompt mode
func initSubscribeFlags(cmd *cobra.Command) {
	cmd.Flags().StringP("prefix", "", "", "subscribe request prefix")
	cmd.Flags().StringArrayVarP(&paths, "path", "", []string{}, "subscribe request paths")
	//cmd.MarkFlagRequired("path")
	cmd.Flags().Uint32P("qos", "q", 0, "qos marking")
	cmd.Flags().BoolP("updates-only", "", false, "only updates to current state should be sent")
	cmd.Flags().StringP("mode", "", "stream", "one of: once, stream, poll")
	cmd.Flags().StringP("stream-mode", "", "target-defined", "one of: on-change, sample, target-defined")
	cmd.Flags().DurationP("sample-interval", "i", 0,
		"sample interval as a decimal number and a suffix unit, such as \"10s\" or \"1m30s\"")
	cmd.Flags().BoolP("suppress-redundant", "", false, "suppress redundant update if the subscribed value didn't not change")
	cmd.Flags().DurationP("heartbeat-interval", "", 0, "heartbeat interval in case suppress-redundant is enabled")
	cmd.Flags().StringSliceP("model", "", []string{}, "subscribe request used model(s)")
	cmd.Flags().Bool("quiet", false, "suppress stdout printing")
	cmd.Flags().StringP("target", "", "", "subscribe request target")
	cmd.Flags().StringSliceP("name", "n", []string{}, "reference subscriptions by name, must be defined in gnmic config file")
	//
	viper.BindPFlag("subscribe-prefix", cmd.LocalFlags().Lookup("prefix"))
	viper.BindPFlag("subscribe-path", cmd.LocalFlags().Lookup("path"))
	viper.BindPFlag("subscribe-qos", cmd.LocalFlags().Lookup("qos"))
	viper.BindPFlag("subscribe-updates-only", cmd.LocalFlags().Lookup("updates-only"))
	viper.BindPFlag("subscribe-mode", cmd.LocalFlags().Lookup("mode"))
	viper.BindPFlag("subscribe-stream-mode", cmd.LocalFlags().Lookup("stream-mode"))
	viper.BindPFlag("subscribe-sample-interval", cmd.LocalFlags().Lookup("sample-interval"))
	viper.BindPFlag("subscribe-suppress-redundant", cmd.LocalFlags().Lookup("suppress-redundant"))
	viper.BindPFlag("subscribe-heartbeat-interval", cmd.LocalFlags().Lookup("heartbeat-interval"))
	viper.BindPFlag("subscribe-sub-model", cmd.LocalFlags().Lookup("model"))
	viper.BindPFlag("subscribe-quiet", cmd.LocalFlags().Lookup("quiet"))
	viper.BindPFlag("subscribe-target", cmd.LocalFlags().Lookup("target"))
	viper.BindPFlag("subscribe-name", cmd.LocalFlags().Lookup("name"))
}

func getOutputs(ctx context.Context) (map[string][]outputs.Output, error) {
	outDef := viper.GetStringMap("outputs")
	if len(outDef) == 0 && !viper.GetBool("quiet") {
		stdoutConfig := map[string]interface{}{
			"type":      "file",
			"file-type": "stdout",
			"format":    viper.GetString("format"),
		}
		outDef["stdout"] = []interface{}{stdoutConfig}
	}
	outputDestinations := make(map[string][]outputs.Output)
	for name, d := range outDef {
		dl := convert(d)
		switch outs := dl.(type) {
		case []interface{}:
			for _, ou := range outs {
				switch ou := ou.(type) {
				case map[string]interface{}:
					if outType, ok := ou["type"]; ok {
						if initalizer, ok := outputs.Outputs[outType.(string)]; ok {
							format, ok := ou["format"]
							if !ok || (ok && format == "") {
								ou["format"] = viper.GetString("format")
							}
							o := initalizer()
							go o.Init(ctx, ou, logger)
							if outputDestinations[name] == nil {
								outputDestinations[name] = make([]outputs.Output, 0)
							}
							outputDestinations[name] = append(outputDestinations[name], o)
							continue
						}
						logger.Printf("unknown output type '%s'", outType)
						continue
					}
					logger.Printf("missing output 'type' under %v", ou)
				default:
					logger.Printf("unknown configuration format expecting a map[string]interface{}: got %T : %v", d, d)
				}
			}
		default:
			return nil, fmt.Errorf("unknown configuration format: %T : %v", d, d)
		}
	}
	return outputDestinations, nil
}

func getSubscriptions() (map[string]*collector.SubscriptionConfig, error) {
	subscriptions := make(map[string]*collector.SubscriptionConfig)
	hi := viper.GetDuration("subscribe-heartbeat-interval")
	si := viper.GetDuration("subscribe-sample-interval")
	var qos *uint32
	// qos value is set to nil by default to enable targets which don't support qos marking
	if viper.IsSet("subscribe-qos") {
		fmt.Println("qos is set")
		q := viper.GetUint32("subscribe-qos")
		qos = &q
	}

	subNames := viper.GetStringSlice("subscribe-name")
	if len(paths) > 0 && len(subNames) > 0 {
		return nil, fmt.Errorf("flags --path and --name cannot be mixed")
	}
	if len(paths) > 0 {
		sub := new(collector.SubscriptionConfig)
		sub.Name = "default"
		sub.Paths = paths
		sub.Prefix = viper.GetString("subscribe-prefix")
		sub.Target = viper.GetString("subscribe-target")
		sub.Mode = viper.GetString("subscribe-mode")
		sub.Encoding = viper.GetString("encoding")
		sub.Qos = qos
		sub.StreamMode = viper.GetString("subscribe-stream-mode")
		sub.HeartbeatInterval = &hi
		sub.SampleInterval = &si
		sub.SuppressRedundant = viper.GetBool("subscribe-suppress-redundant")
		sub.UpdatesOnly = viper.GetBool("subscribe-updates-only")
		sub.Models = viper.GetStringSlice("models")
		subscriptions["default"] = sub
		return subscriptions, nil
	}
	subDef := viper.GetStringMap("subscriptions")
	if viper.GetBool("debug") {
		logger.Printf("subscription map: %v+", subDef)
	}
	for sn, s := range subDef {
		sub := new(collector.SubscriptionConfig)
		decoder, err := mapstructure.NewDecoder(
			&mapstructure.DecoderConfig{
				DecodeHook: mapstructure.StringToTimeDurationHookFunc(),
				Result:     sub,
			})
		if err != nil {
			return nil, err
		}
		err = decoder.Decode(s)
		if err != nil {
			return nil, err
		}
		sub.Name = sn

		// inherit global "subscribe-*" option if it's not set
		if sub.SampleInterval == nil {
			sub.SampleInterval = &si
		}
		if sub.HeartbeatInterval == nil {
			sub.HeartbeatInterval = &hi
		}
		if sub.Encoding == "" {
			sub.Encoding = viper.GetString("encoding")
		}
		if sub.Mode == "" {
			sub.Mode = viper.GetString("subscribe-mode")
		}
		if strings.ToUpper(sub.Mode) == "STREAM" && sub.StreamMode == "" {
			sub.StreamMode = viper.GetString("subscribe-stream-mode")
		}
		if sub.Qos == nil {
			sub.Qos = qos
		}
		subscriptions[sn] = sub
	}
	if len(subNames) == 0 {
		return subscriptions, nil
	}
	filteredSubscriptions := make(map[string]*collector.SubscriptionConfig)
	notFound := make([]string, 0)
	for _, name := range subNames {
		if s, ok := subscriptions[name]; ok {
			filteredSubscriptions[name] = s
		} else {
			notFound = append(notFound, name)
		}
	}
	if len(notFound) > 0 {
		return nil, fmt.Errorf("named subscription(s) not found in config file: %v", notFound)
	}
	return filteredSubscriptions, nil
}
<commit_msg>fixed typo<commit_after>// Copyright © 2020 Karim Radhouani <medkarimrdi@gmail.com>
//
// 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 cmd

import (
	"context"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	_ "github.com/karimra/gnmic/outputs/all"
	"github.com/manifoldco/promptui"
	"github.com/mitchellh/mapstructure"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

const defaultRetryTimer = 10 * time.Second

var subscriptionModes = [][2]string{
	{"once", "a single request/response channel. The target creates the relevant update messages, transmits them, and subsequently closes the RPC"},
	{"stream", "long-lived subscriptions which continue to transmit updates relating to the set of paths that are covered within the subscription indefinitely"},
	{"poll", "on-demand retrieval of data items via long-lived RPCs"},
}
var streamSubscriptionModes = [][2]string{
	{"target-defined", "the target MUST determine the best type of subscription to be created on a per-leaf basis"},
	{"sample", "the value of the data item(s) MUST be sent once per sample interval to the client"},
	{"on-change", "data updates are only sent when the value of the data item changes"},
}

// subscribeCmd represents the subscribe command
var subscribeCmd = &cobra.Command{
	Use:     "subscribe",
	Aliases: []string{"sub"},
	Short:   "subscribe to gnmi updates on targets",
	Annotations: map[string]string{
		"--path":        "XPATH",
		"--prefix":      "PREFIX",
		"--model":       "MODEL",
		"--mode":        "SUBSC_MODE",
		"--stream-mode": "STREAM_MODE",
	},

	RunE: func(cmd *cobra.Command, args []string) error {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		setupCloseHandler(cancel)
		debug := viper.GetBool("debug")
		targetsConfig, err := createTargets()
		if err != nil {
			return fmt.Errorf("failed getting targets config: %v", err)
		}
		if debug {
			logger.Printf("targets: %s", targetsConfig)
		}
		subscriptionsConfig, err := getSubscriptions()
		if err != nil {
			return fmt.Errorf("failed getting subscriptions config: %v", err)
		}
		if debug {
			logger.Printf("subscriptions: %s", subscriptionsConfig)
		}
		outs, err := getOutputs(ctx)
		if err != nil {
			return err
		}
		if debug {
			logger.Printf("outputs: %+v", outs)
		}

		cfg := &collector.Config{
			PrometheusAddress:   viper.GetString("prometheus-address"),
			Debug:               viper.GetBool("debug"),
			Format:              viper.GetString("format"),
			TargetReceiveBuffer: viper.GetUint("target-buffer-size"),
			RetryTimer:          viper.GetDuration("retry-timer"),
		}

		coll := collector.NewCollector(cfg, targetsConfig, subscriptionsConfig, outs, createCollectorDialOpts(), logger)

		wg := new(sync.WaitGroup)
		wg.Add(len(coll.Targets))
		for name := range coll.Targets {
			go func(tn string) {
				defer wg.Done()
				tRetryTimer := coll.Targets[tn].Config.RetryTimer
				for {
					err = coll.Subscribe(ctx, tn)
					if err != nil {
						if errors.Is(err, context.DeadlineExceeded) {
							logger.Printf("failed to initialize target '%s' timeout (%s) reached", tn, targetsConfig[tn].Timeout)
						} else {
							logger.Printf("failed to initialize target '%s': %v", tn, err)
						}
						logger.Printf("retrying target %s in %s", tn, tRetryTimer)
						time.Sleep(tRetryTimer)
						continue
					}
					return
				}
			}(name)
		}
		wg.Wait()
		polledTargetsSubscriptions := coll.PolledSubscriptionsTargets()
		if len(polledTargetsSubscriptions) > 0 {
			pollTargets := make([]string, 0, len(polledTargetsSubscriptions))
			for t := range polledTargetsSubscriptions {
				pollTargets = append(pollTargets, t)
			}
			sort.Slice(pollTargets, func(i, j int) bool {
				return pollTargets[i] < pollTargets[j]
			})
			s := promptui.Select{
				Label:        "select target to poll",
				Items:        pollTargets,
				HideSelected: true,
			}
			waitChan := make(chan struct{}, 1)
			waitChan <- struct{}{}
			mo := &collector.MarshalOptions{
				Multiline: true,
				Indent:    "  ",
				Format:    viper.GetString("format")}
			go func() {
				for {
					select {
					case <-waitChan:
						_, name, err := s.Run()
						if err != nil {
							fmt.Printf("failed selecting target to poll: %v\n", err)
							continue
						}
						ss := promptui.Select{
							Label:        "select subscription to poll",
							Items:        polledTargetsSubscriptions[name],
							HideSelected: true,
						}
						_, subName, err := ss.Run()
						if err != nil {
							fmt.Printf("failed selecting subscription to poll: %v\n", err)
							continue
						}
						response, err := coll.TargetPoll(name, subName)
						if err != nil && err != io.EOF {
							fmt.Printf("target '%s', subscription '%s': poll response error:%v\n", name, subName, err)
							continue
						}
						if response == nil {
							fmt.Printf("received empty response from target '%s'\n", name)
							continue
						}
						switch rsp := response.Response.(type) {
						case *gnmi.SubscribeResponse_SyncResponse:
							fmt.Printf("received sync response '%t' from '%s'\n", rsp.SyncResponse, name)
							waitChan <- struct{}{}
							continue
						}
						b, err := mo.Marshal(response, nil)
						if err != nil {
							fmt.Printf("target '%s', subscription '%s': poll response formatting error:%v\n", name, subName, err)
							fmt.Println(string(b))
							waitChan <- struct{}{}
							continue
						}
						fmt.Println(string(b))
						waitChan <- struct{}{}
					case <-ctx.Done():
						return
					}
				}
			}()
		}
		coll.Start(ctx)
		return nil
	},
	PostRun: func(cmd *cobra.Command, args []string) {
		cmd.ResetFlags()
		initSubscribeFlags(cmd)
	},
	SilenceUsage: true,
}

func init() {
	rootCmd.AddCommand(subscribeCmd)
	initSubscribeFlags(subscribeCmd)
}

// used to init or reset subscribeCmd flags for gnmic-prompt mode
func initSubscribeFlags(cmd *cobra.Command) {
	cmd.Flags().StringP("prefix", "", "", "subscribe request prefix")
	cmd.Flags().StringArrayVarP(&paths, "path", "", []string{}, "subscribe request paths")
	//cmd.MarkFlagRequired("path")
	cmd.Flags().Uint32P("qos", "q", 0, "qos marking")
	cmd.Flags().BoolP("updates-only", "", false, "only updates to current state should be sent")
	cmd.Flags().StringP("mode", "", "stream", "one of: once, stream, poll")
	cmd.Flags().StringP("stream-mode", "", "target-defined", "one of: on-change, sample, target-defined")
	cmd.Flags().DurationP("sample-interval", "i", 0,
		"sample interval as a decimal number and a suffix unit, such as \"10s\" or \"1m30s\"")
	cmd.Flags().BoolP("suppress-redundant", "", false, "suppress redundant update if the subscribed value didn't not change")
	cmd.Flags().DurationP("heartbeat-interval", "", 0, "heartbeat interval in case suppress-redundant is enabled")
	cmd.Flags().StringSliceP("model", "", []string{}, "subscribe request used model(s)")
	cmd.Flags().Bool("quiet", false, "suppress stdout printing")
	cmd.Flags().StringP("target", "", "", "subscribe request target")
	cmd.Flags().StringSliceP("name", "n", []string{}, "reference subscriptions by name, must be defined in gnmic config file")
	//
	viper.BindPFlag("subscribe-prefix", cmd.LocalFlags().Lookup("prefix"))
	viper.BindPFlag("subscribe-path", cmd.LocalFlags().Lookup("path"))
	viper.BindPFlag("subscribe-qos", cmd.LocalFlags().Lookup("qos"))
	viper.BindPFlag("subscribe-updates-only", cmd.LocalFlags().Lookup("updates-only"))
	viper.BindPFlag("subscribe-mode", cmd.LocalFlags().Lookup("mode"))
	viper.BindPFlag("subscribe-stream-mode", cmd.LocalFlags().Lookup("stream-mode"))
	viper.BindPFlag("subscribe-sample-interval", cmd.LocalFlags().Lookup("sample-interval"))
	viper.BindPFlag("subscribe-suppress-redundant", cmd.LocalFlags().Lookup("suppress-redundant"))
	viper.BindPFlag("subscribe-heartbeat-interval", cmd.LocalFlags().Lookup("heartbeat-interval"))
	viper.BindPFlag("subscribe-sub-model", cmd.LocalFlags().Lookup("model"))
	viper.BindPFlag("subscribe-quiet", cmd.LocalFlags().Lookup("quiet"))
	viper.BindPFlag("subscribe-target", cmd.LocalFlags().Lookup("target"))
	viper.BindPFlag("subscribe-name", cmd.LocalFlags().Lookup("name"))
}

func getOutputs(ctx context.Context) (map[string][]outputs.Output, error) {
	outDef := viper.GetStringMap("outputs")
	if len(outDef) == 0 && !viper.GetBool("quiet") {
		stdoutConfig := map[string]interface{}{
			"type":      "file",
			"file-type": "stdout",
			"format":    viper.GetString("format"),
		}
		outDef["stdout"] = []interface{}{stdoutConfig}
	}
	outputDestinations := make(map[string][]outputs.Output)
	for name, d := range outDef {
		dl := convert(d)
		switch outs := dl.(type) {
		case []interface{}:
			for _, ou := range outs {
				switch ou := ou.(type) {
				case map[string]interface{}:
					if outType, ok := ou["type"]; ok {
						if initializer, ok := outputs.Outputs[outType.(string)]; ok {
							format, ok := ou["format"]
							if !ok || (ok && format == "") {
								ou["format"] = viper.GetString("format")
							}
							o := initializer()
							go o.Init(ctx, ou, logger)
							if outputDestinations[name] == nil {
								outputDestinations[name] = make([]outputs.Output, 0)
							}
							outputDestinations[name] = append(outputDestinations[name], o)
							continue
						}
						logger.Printf("unknown output type '%s'", outType)
						continue
					}
					logger.Printf("missing output 'type' under %v", ou)
				default:
					logger.Printf("unknown configuration format expecting a map[string]interface{}: got %T : %v", d, d)
				}
			}
		default:
			return nil, fmt.Errorf("unknown configuration format: %T : %v", d, d)
		}
	}
	return outputDestinations, nil
}

func getSubscriptions() (map[string]*collector.SubscriptionConfig, error) {
	subscriptions := make(map[string]*collector.SubscriptionConfig)
	hi := viper.GetDuration("subscribe-heartbeat-interval")
	si := viper.GetDuration("subscribe-sample-interval")
	var qos *uint32
	// qos value is set to nil by default to enable targets which don't support qos marking
	if viper.IsSet("subscribe-qos") {
		fmt.Println("qos is set")
		q := viper.GetUint32("subscribe-qos")
		qos = &q
	}

	subNames := viper.GetStringSlice("subscribe-name")
	if len(paths) > 0 && len(subNames) > 0 {
		return nil, fmt.Errorf("flags --path and --name cannot be mixed")
	}
	if len(paths) > 0 {
		sub := new(collector.SubscriptionConfig)
		sub.Name = "default"
		sub.Paths = paths
		sub.Prefix = viper.GetString("subscribe-prefix")
		sub.Target = viper.GetString("subscribe-target")
		sub.Mode = viper.GetString("subscribe-mode")
		sub.Encoding = viper.GetString("encoding")
		sub.Qos = qos
		sub.StreamMode = viper.GetString("subscribe-stream-mode")
		sub.HeartbeatInterval = &hi
		sub.SampleInterval = &si
		sub.SuppressRedundant = viper.GetBool("subscribe-suppress-redundant")
		sub.UpdatesOnly = viper.GetBool("subscribe-updates-only")
		sub.Models = viper.GetStringSlice("models")
		subscriptions["default"] = sub
		return subscriptions, nil
	}
	subDef := viper.GetStringMap("subscriptions")
	if viper.GetBool("debug") {
		logger.Printf("subscription map: %v+", subDef)
	}
	for sn, s := range subDef {
		sub := new(collector.SubscriptionConfig)
		decoder, err := mapstructure.NewDecoder(
			&mapstructure.DecoderConfig{
				DecodeHook: mapstructure.StringToTimeDurationHookFunc(),
				Result:     sub,
			})
		if err != nil {
			return nil, err
		}
		err = decoder.Decode(s)
		if err != nil {
			return nil, err
		}
		sub.Name = sn

		// inherit global "subscribe-*" option if it's not set
		if sub.SampleInterval == nil {
			sub.SampleInterval = &si
		}
		if sub.HeartbeatInterval == nil {
			sub.HeartbeatInterval = &hi
		}
		if sub.Encoding == "" {
			sub.Encoding = viper.GetString("encoding")
		}
		if sub.Mode == "" {
			sub.Mode = viper.GetString("subscribe-mode")
		}
		if strings.ToUpper(sub.Mode) == "STREAM" && sub.StreamMode == "" {
			sub.StreamMode = viper.GetString("subscribe-stream-mode")
		}
		if sub.Qos == nil {
			sub.Qos = qos
		}
		subscriptions[sn] = sub
	}
	if len(subNames) == 0 {
		return subscriptions, nil
	}
	filteredSubscriptions := make(map[string]*collector.SubscriptionConfig)
	notFound := make([]string, 0)
	for _, name := range subNames {
		if s, ok := subscriptions[name]; ok {
			filteredSubscriptions[name] = s
		} else {
			notFound = append(notFound, name)
		}
	}
	if len(notFound) > 0 {
		return nil, fmt.Errorf("named subscription(s) not found in config file: %v", notFound)
	}
	return filteredSubscriptions, nil
}
<|endoftext|>


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

package main

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
)

type talksBuilder struct {
}

func (b talksBuilder) Signature(heads map[string]string) string {
	return heads["talks"]
}

const talksToolsRev = "1f1b3322f67af76803c942fd237291538ec68262"

func (b talksBuilder) Init(dir, hostport string, heads map[string]string) (*exec.Cmd, error) {
	toolsDir := filepath.Join(dir, "gopath/src/golang.org/x/tools")
	if err := checkout(repoURL+"tools", talksToolsRev, toolsDir); err != nil {
		return nil, err
	}
	talksDir := filepath.Join(dir, "gopath/src/golang.org/x/talks")
	if err := checkout(repoURL+"talks", heads["talks"], talksDir); err != nil {
		return nil, err
	}

	goDir := os.Getenv("GOROOT_BOOTSTRAP")
	if goDir == "" {
		goDir = runtime.GOROOT()
	}
	goBin := filepath.Join(goDir, "bin/go")
	goPath := filepath.Join(dir, "gopath")
	presentPath := "golang.org/x/tools/cmd/present"
	install := exec.Command(goBin, "install", "-tags=appenginevm", presentPath)
	install.Env = []string{"GOROOT=" + goDir, "GOPATH=" + goPath}
	if err := runErr(install); err != nil {
		return nil, err
	}

	talksBin := filepath.Join(goPath, "bin/present")
	presentSrc := filepath.Join(goPath, "src", presentPath)
	present := exec.Command(talksBin, "-http="+hostport, "-base="+presentSrc)
	present.Dir = talksDir
	// TODO(adg): log this somewhere useful
	present.Stdout = os.Stdout
	present.Stderr = os.Stderr
	if err := present.Start(); err != nil {
		return nil, err
	}
	return present, nil
}

var talksMsg = []byte("Talks - The Go Programming Language")

func (b talksBuilder) HealthCheck(hostport string) error {
	body, err := getOK(fmt.Sprintf("http://%v/", hostport))
	if err != nil {
		return err
	}
	if !bytes.Contains(body, talksMsg) {
		return errors.New("couldn't match string")
	}
	return nil
}
<commit_msg>cmd/tip: update talks rev to pick up CL 33578<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
)

type talksBuilder struct {
}

func (b talksBuilder) Signature(heads map[string]string) string {
	return heads["talks"]
}

const talksToolsRev = "f19f04f5492745c70cc5906b15b125b7a3d3b1a3"

func (b talksBuilder) Init(dir, hostport string, heads map[string]string) (*exec.Cmd, error) {
	toolsDir := filepath.Join(dir, "gopath/src/golang.org/x/tools")
	if err := checkout(repoURL+"tools", talksToolsRev, toolsDir); err != nil {
		return nil, err
	}
	talksDir := filepath.Join(dir, "gopath/src/golang.org/x/talks")
	if err := checkout(repoURL+"talks", heads["talks"], talksDir); err != nil {
		return nil, err
	}

	goDir := os.Getenv("GOROOT_BOOTSTRAP")
	if goDir == "" {
		goDir = runtime.GOROOT()
	}
	goBin := filepath.Join(goDir, "bin/go")
	goPath := filepath.Join(dir, "gopath")
	presentPath := "golang.org/x/tools/cmd/present"
	install := exec.Command(goBin, "install", "-tags=appenginevm", presentPath)
	install.Env = []string{"GOROOT=" + goDir, "GOPATH=" + goPath}
	if err := runErr(install); err != nil {
		return nil, err
	}

	talksBin := filepath.Join(goPath, "bin/present")
	presentSrc := filepath.Join(goPath, "src", presentPath)
	present := exec.Command(talksBin, "-http="+hostport, "-base="+presentSrc)
	present.Dir = talksDir
	// TODO(adg): log this somewhere useful
	present.Stdout = os.Stdout
	present.Stderr = os.Stderr
	if err := present.Start(); err != nil {
		return nil, err
	}
	return present, nil
}

var talksMsg = []byte("Talks - The Go Programming Language")

func (b talksBuilder) HealthCheck(hostport string) error {
	body, err := getOK(fmt.Sprintf("http://%v/", hostport))
	if err != nil {
		return err
	}
	if !bytes.Contains(body, talksMsg) {
		return errors.New("couldn't match string")
	}
	return nil
}
<|endoftext|>


<commit_before>package integration_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os/exec"

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

var _ = Describe("Integration", func() {
	var (
		session    *gexec.Session
		address    string
		listenPort int
	)

	var serverIsAvailable = func() error {
		return VerifyTCPConnection(address)
	}

	BeforeEach(func() {
		listenPort = rand.Intn(1000) + 5000
		address = fmt.Sprintf("127.0.0.1:%d", listenPort)

		exampleAppCmd := exec.Command(exampleAppPath)
		exampleAppCmd.Env = []string{
			fmt.Sprintf("PORT=%d", listenPort),
		}
		var err error
		session, err = gexec.Start(exampleAppCmd, GinkgoWriter, GinkgoWriter)
		Expect(err).NotTo(HaveOccurred())

		Eventually(serverIsAvailable, DEFAULT_TIMEOUT).Should(Succeed())
	})

	AfterEach(func() {
		session.Interrupt()
		Eventually(session, DEFAULT_TIMEOUT).Should(gexec.Exit())
	})
	Describe("boring server behavior", func() {
		It("should boot and gracefully terminate", func() {
			Consistently(session).ShouldNot(gexec.Exit())

			session.Interrupt()
			Eventually(session, DEFAULT_TIMEOUT).Should(gexec.Exit())
		})
	})

	Describe("endpoints", func() {
		It("should respond to GET / with info", func() {
			response, err := http.DefaultClient.Get("http://" + address + "/")
			Expect(err).NotTo(HaveOccurred())
			defer response.Body.Close()
			Expect(response.StatusCode).To(Equal(200))

			responseBytes, err := ioutil.ReadAll(response.Body)
			Expect(err).NotTo(HaveOccurred())

			var responseData struct {
				ListenAddresses []string
				Port            int
			}

			Expect(json.Unmarshal(responseBytes, &responseData)).To(Succeed())

			Expect(responseData.ListenAddresses).To(ContainElement("127.0.0.1"))
			Expect(responseData.Port).To(Equal(listenPort))
		})

		It("should respond to /proxy by proxying the request to the provided address", func() {
			response, err := http.DefaultClient.Get("http://" + address + "/proxy/example.com")
			Expect(err).NotTo(HaveOccurred())
			defer response.Body.Close()
			Expect(response.StatusCode).To(Equal(200))

			responseBytes, err := ioutil.ReadAll(response.Body)
			Expect(err).NotTo(HaveOccurred())
			Expect(responseBytes).To(ContainSubstring("Example Domain"))
		})

		It("should report latency stats on /stats", func() {
			response, err := http.DefaultClient.Get("http://" + address + "/proxy/example.com")
			Expect(err).NotTo(HaveOccurred())
			Expect(response.StatusCode).To(Equal(200))

			statsResponse, err := http.DefaultClient.Get("http://" + address + "/stats")
			Expect(err).NotTo(HaveOccurred())
			defer statsResponse.Body.Close()

			responseBytes, err := ioutil.ReadAll(statsResponse.Body)
			Expect(err).NotTo(HaveOccurred())
			var statsJSON struct {
				Latency []float64
			}
			Expect(json.Unmarshal(responseBytes, &statsJSON)).To(Succeed())
			Expect(len(statsJSON.Latency)).To(BeNumerically(">=", 1))
		})

		Context("when the proxy destination is invalid", func() {
			It("logs the error", func() {
				response, err := http.DefaultClient.Get("http://" + address + "/proxy/////!!")
				Expect(err).NotTo(HaveOccurred())
				defer response.Body.Close()
				Expect(response.StatusCode).To(Equal(500))

				Expect(session.Err.Contents()).To(ContainSubstring("request failed: Get"))
			})
		})
	})
})
<commit_msg>Fix proxy integration test to wait for request failure<commit_after>package integration_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"os/exec"

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

var _ = Describe("Integration", func() {
	var (
		session    *gexec.Session
		address    string
		listenPort int
	)

	var serverIsAvailable = func() error {
		return VerifyTCPConnection(address)
	}

	BeforeEach(func() {
		listenPort = rand.Intn(1000) + 5000
		address = fmt.Sprintf("127.0.0.1:%d", listenPort)

		exampleAppCmd := exec.Command(exampleAppPath)
		exampleAppCmd.Env = []string{
			fmt.Sprintf("PORT=%d", listenPort),
		}
		var err error
		session, err = gexec.Start(exampleAppCmd, GinkgoWriter, GinkgoWriter)
		Expect(err).NotTo(HaveOccurred())

		Eventually(serverIsAvailable, DEFAULT_TIMEOUT).Should(Succeed())
	})

	AfterEach(func() {
		session.Interrupt()
		Eventually(session, DEFAULT_TIMEOUT).Should(gexec.Exit())
	})
	Describe("boring server behavior", func() {
		It("should boot and gracefully terminate", func() {
			Consistently(session).ShouldNot(gexec.Exit())

			session.Interrupt()
			Eventually(session, DEFAULT_TIMEOUT).Should(gexec.Exit())
		})
	})

	Describe("endpoints", func() {
		It("should respond to GET / with info", func() {
			response, err := http.DefaultClient.Get("http://" + address + "/")
			Expect(err).NotTo(HaveOccurred())
			defer response.Body.Close()
			Expect(response.StatusCode).To(Equal(200))

			responseBytes, err := ioutil.ReadAll(response.Body)
			Expect(err).NotTo(HaveOccurred())

			var responseData struct {
				ListenAddresses []string
				Port            int
			}

			Expect(json.Unmarshal(responseBytes, &responseData)).To(Succeed())

			Expect(responseData.ListenAddresses).To(ContainElement("127.0.0.1"))
			Expect(responseData.Port).To(Equal(listenPort))
		})

		It("should respond to /proxy by proxying the request to the provided address", func() {
			response, err := http.DefaultClient.Get("http://" + address + "/proxy/example.com")
			Expect(err).NotTo(HaveOccurred())
			defer response.Body.Close()
			Expect(response.StatusCode).To(Equal(200))

			responseBytes, err := ioutil.ReadAll(response.Body)
			Expect(err).NotTo(HaveOccurred())
			Expect(responseBytes).To(ContainSubstring("Example Domain"))
		})

		It("should report latency stats on /stats", func() {
			response, err := http.DefaultClient.Get("http://" + address + "/proxy/example.com")
			Expect(err).NotTo(HaveOccurred())
			Expect(response.StatusCode).To(Equal(200))

			statsResponse, err := http.DefaultClient.Get("http://" + address + "/stats")
			Expect(err).NotTo(HaveOccurred())
			defer statsResponse.Body.Close()

			responseBytes, err := ioutil.ReadAll(statsResponse.Body)
			Expect(err).NotTo(HaveOccurred())
			var statsJSON struct {
				Latency []float64
			}
			Expect(json.Unmarshal(responseBytes, &statsJSON)).To(Succeed())
			Expect(len(statsJSON.Latency)).To(BeNumerically(">=", 1))
		})

		Context("when the proxy destination is invalid", func() {
			It("logs the error", func() {
				response, err := http.DefaultClient.Get("http://" + address + "/proxy/////!!")
				Expect(err).NotTo(HaveOccurred())
				defer response.Body.Close()
				Expect(response.StatusCode).To(Equal(500))

				Eventually(session.Err.Contents).Should(ContainSubstring("request failed: Get"))
			})
		})
	})
})
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"log"
	"math"
	"strconv"

	"github.com/turing-complete/laboratory/src/internal/command"
	"github.com/turing-complete/laboratory/src/internal/config"
	"github.com/turing-complete/laboratory/src/internal/database"
	"github.com/turing-complete/laboratory/src/internal/quantity"
	"github.com/turing-complete/laboratory/src/internal/solution"
	"github.com/turing-complete/laboratory/src/internal/support"
	"github.com/turing-complete/laboratory/src/internal/system"
	"github.com/turing-complete/laboratory/src/internal/uncertainty"
)

const (
	maxSteps = 20
)

var (
	approximateFile = flag.String("approximate", "", "an output of `approximate` (required)")
	outputFile      = flag.String("o", "", "an output file (required)")
	sampleSeed      = flag.String("s", "", "a seed for generating samples")
	sampleCount     = flag.String("n", "", "the number of samples")
)

type Config *config.Assessment

func main() {
	command.Run(function)
}

func function(config *config.Config) error {
	if len(*sampleSeed) > 0 {
		if number, err := strconv.ParseInt(*sampleSeed, 0, 64); err != nil {
			return err
		} else {
			config.Assessment.Seed = number
		}
	}
	if len(*sampleCount) > 0 {
		if number, err := strconv.ParseUint(*sampleCount, 0, 64); err != nil {
			return err
		} else {
			config.Assessment.Samples = uint(number)
		}
	}

	if config.Assessment.Samples == 0 {
		return errors.New("the number of samples should be positive")
	}

	approximate, err := database.Open(*approximateFile)
	if err != nil {
		return err
	}
	defer approximate.Close()

	output, err := database.Create(*outputFile)
	if err != nil {
		return err
	}
	defer output.Close()

	system, err := system.New(&config.System)
	if err != nil {
		return err
	}

	euncertainty, err := uncertainty.NewEpistemic(system, &config.Uncertainty)
	if err != nil {
		return err
	}
	equantity, err := quantity.New(system, euncertainty, &config.Quantity)
	if err != nil {
		return err
	}

	auncertainty, err := uncertainty.NewAleatory(system, &config.Uncertainty)
	if err != nil {
		return err
	}
	aquantity, err := quantity.New(system, auncertainty, &config.Quantity)
	if err != nil {
		return err
	}

	ni, no := equantity.Dimensions()

	surrogate := new(solution.Surrogate)
	if err = approximate.Get("surrogate", surrogate); err != nil {
		return err
	}

	solution, err := solution.New(ni, no, &config.Solution)
	if err != nil {
		return err
	}

	ns := config.Assessment.Samples

	epoints, apoints := generate(equantity, aquantity, ns, config.Assessment.Seed)

	log.Printf("Evaluating the surrogate model at %d points...\n", ns)
	log.Printf("%5s %15s\n", "Step", "Nodes")

	nk := uint(len(surrogate.Active))

	steps := make([]uint, nk)
	values := make([]float64, 0, ns*no)

	k, Δ := uint(0), float64(nk-1)/(math.Min(maxSteps, float64(nk))-1)

	for i, na := uint(0), uint(0); i < nk; i++ {
		na += surrogate.Active[i]
		steps[k] += surrogate.Active[i]

		if i != uint(float64(k)*Δ+0.5) {
			continue
		}
		k++

		log.Printf("%5d %15d\n", i, na)

		s := *surrogate
		s.Nodes = na
		s.Indices = s.Indices[:na*ni]
		s.Surpluses = s.Surpluses[:na*no]

		if !solution.Validate(&s) {
			panic("something went wrong")
		}

		values = append(values, solution.Evaluate(&s, epoints)...)
	}

	nk, steps = k, steps[:k]

	log.Println("Done.")

	if err := output.Put("surrogate", *surrogate); err != nil {
		return err
	}
	if err := output.Put("points", apoints, ni, ns); err != nil {
		return err
	}
	if err := output.Put("steps", steps); err != nil {
		return err
	}
	if err := output.Put("values", values, no, ns, nk); err != nil {
		return err
	}

	return nil
}

func generate(into, from quantity.Quantity, ns uint, seed int64) ([]float64, []float64) {
	nii, _ := into.Dimensions()
	nif, _ := from.Dimensions()

	zi := make([]float64, nii*ns)
	zf := support.Generate(nif, ns, seed)

	for i := uint(0); i < ns; i++ {
		copy(zi[i*nii:(i+1)*nii], into.Forward(from.Backward(zf[i*nif:(i+1)*nif])))
	}

	return zi, zf
}
<commit_msg>s/predict: fix the dimensions of points<commit_after>package main

import (
	"errors"
	"flag"
	"log"
	"math"
	"strconv"

	"github.com/turing-complete/laboratory/src/internal/command"
	"github.com/turing-complete/laboratory/src/internal/config"
	"github.com/turing-complete/laboratory/src/internal/database"
	"github.com/turing-complete/laboratory/src/internal/quantity"
	"github.com/turing-complete/laboratory/src/internal/solution"
	"github.com/turing-complete/laboratory/src/internal/support"
	"github.com/turing-complete/laboratory/src/internal/system"
	"github.com/turing-complete/laboratory/src/internal/uncertainty"
)

const (
	maxSteps = 20
)

var (
	approximateFile = flag.String("approximate", "", "an output of `approximate` (required)")
	outputFile      = flag.String("o", "", "an output file (required)")
	sampleSeed      = flag.String("s", "", "a seed for generating samples")
	sampleCount     = flag.String("n", "", "the number of samples")
)

type Config *config.Assessment

func main() {
	command.Run(function)
}

func function(config *config.Config) error {
	if len(*sampleSeed) > 0 {
		if number, err := strconv.ParseInt(*sampleSeed, 0, 64); err != nil {
			return err
		} else {
			config.Assessment.Seed = number
		}
	}
	if len(*sampleCount) > 0 {
		if number, err := strconv.ParseUint(*sampleCount, 0, 64); err != nil {
			return err
		} else {
			config.Assessment.Samples = uint(number)
		}
	}

	if config.Assessment.Samples == 0 {
		return errors.New("the number of samples should be positive")
	}

	approximate, err := database.Open(*approximateFile)
	if err != nil {
		return err
	}
	defer approximate.Close()

	output, err := database.Create(*outputFile)
	if err != nil {
		return err
	}
	defer output.Close()

	system, err := system.New(&config.System)
	if err != nil {
		return err
	}

	euncertainty, err := uncertainty.NewEpistemic(system, &config.Uncertainty)
	if err != nil {
		return err
	}
	equantity, err := quantity.New(system, euncertainty, &config.Quantity)
	if err != nil {
		return err
	}

	auncertainty, err := uncertainty.NewAleatory(system, &config.Uncertainty)
	if err != nil {
		return err
	}
	aquantity, err := quantity.New(system, auncertainty, &config.Quantity)
	if err != nil {
		return err
	}

	nie, no := equantity.Dimensions()
	nia, _ := aquantity.Dimensions()

	surrogate := new(solution.Surrogate)
	if err = approximate.Get("surrogate", surrogate); err != nil {
		return err
	}

	solution, err := solution.New(nie, no, &config.Solution)
	if err != nil {
		return err
	}

	ns := config.Assessment.Samples

	epoints, apoints := generate(equantity, aquantity, ns, config.Assessment.Seed)

	log.Printf("Evaluating the surrogate model at %d points...\n", ns)
	log.Printf("%5s %15s\n", "Step", "Nodes")

	nk := uint(len(surrogate.Active))

	steps := make([]uint, nk)
	values := make([]float64, 0, ns*no)

	k, Δ := uint(0), float64(nk-1)/(math.Min(maxSteps, float64(nk))-1)

	for i, na := uint(0), uint(0); i < nk; i++ {
		na += surrogate.Active[i]
		steps[k] += surrogate.Active[i]

		if i != uint(float64(k)*Δ+0.5) {
			continue
		}
		k++

		log.Printf("%5d %15d\n", i, na)

		s := *surrogate
		s.Nodes = na
		s.Indices = s.Indices[:na*nie]
		s.Surpluses = s.Surpluses[:na*no]

		if !solution.Validate(&s) {
			panic("something went wrong")
		}

		values = append(values, solution.Evaluate(&s, epoints)...)
	}

	nk, steps = k, steps[:k]

	log.Println("Done.")

	if err := output.Put("surrogate", *surrogate); err != nil {
		return err
	}
	if err := output.Put("points", apoints, nia, ns); err != nil {
		return err
	}
	if err := output.Put("steps", steps); err != nil {
		return err
	}
	if err := output.Put("values", values, no, ns, nk); err != nil {
		return err
	}

	return nil
}

func generate(into, from quantity.Quantity, ns uint, seed int64) ([]float64, []float64) {
	nii, _ := into.Dimensions()
	nif, _ := from.Dimensions()

	zi := make([]float64, nii*ns)
	zf := support.Generate(nif, ns, seed)

	for i := uint(0); i < ns; i++ {
		copy(zi[i*nii:(i+1)*nii], into.Forward(from.Backward(zf[i*nif:(i+1)*nif])))
	}

	return zi, zf
}
<|endoftext|>


<commit_before>// +build !android

package osversion

import (
	"C"
	"errors"
	"fmt"
	"io/ioutil"
	"regexp"
	"syscall"
)

func GetString() (string, error) {
	var uts syscall.Utsname
	err := syscall.Uname(&uts)
	if err != nil {
		return "", errors.New(fmt.Sprintf("Error calling system function 'uname': %s", err))
	}
	return fmt.Sprintf("%s", int8SliceToString(uts.Release[:])), nil
}

func GetHumanReadable() (string, error) {
	// Kernel version
	kernel, err := GetString()
	if err != nil {
		return "", err
	}

	// Try to get the distribution info
	fData, err := ioutil.ReadFile("/etc/os-release")
	if err != nil {
		return fmt.Sprintf("kernel: %s", kernel), nil
	}

	// At least Fedora, Debian, Ubuntu and Arch support this approach
	// and provide the PRETTY_NAME field
	reg1 := regexp.MustCompile("PRETTY_NAME=\".+\"")
	reg2 := regexp.MustCompile("\".+\"")
	dstrBytes := reg2.Find(reg1.Find(fData))
	distribution := string(dstrBytes[1 : len(dstrBytes)-1])

	return fmt.Sprintf("%s kernel: %s", distribution, kernel), nil
}

func int8SliceToString(ca []int8) string {
	s := make([]byte, len(ca))
	strpos := 0
	for strpos < len(ca) {
		if ca[strpos] == 0 {
			break
		}
		s[strpos] = uint8(ca[strpos])
		strpos++
	}
	return string(s[:strpos])
}
<commit_msg>Don't use a function so casting can be applied to each array element directly (bypass Go's type checking for slices)<commit_after>// +build !android

package osversion

import (
	"C"
	"errors"
	"fmt"
	"io/ioutil"
	"regexp"
	"syscall"
)

func GetString() (string, error) {
	var uts syscall.Utsname
	err := syscall.Uname(&uts)
	if err != nil {
		return "", errors.New(fmt.Sprintf("Error calling system function 'uname': %s", err))
	}

	// Due to a mismatch in the uts.Release types depending on the architecture, we are
	// forced to implement it right here to bypass Go's type checking of slices
	utsRelease := uts.Release[:]
	s := make([]byte, len(utsRelease))
	strpos := 0
	for strpos < len(utsRelease) {
		if utsRelease[strpos] == 0 {
			break
		}
		s[strpos] = uint8(utsRelease[strpos])
		strpos++
	}

	return fmt.Sprintf("%s", string(s[:strpos])), nil
}

func GetHumanReadable() (string, error) {
	// Kernel version
	kernel, err := GetString()
	if err != nil {
		return "", err
	}

	// Try to get the distribution info
	fData, err := ioutil.ReadFile("/etc/os-release")
	if err != nil {
		return fmt.Sprintf("kernel: %s", kernel), nil
	}

	// At least Fedora, Debian, Ubuntu and Arch support this approach
	// and provide the PRETTY_NAME field
	reg1 := regexp.MustCompile("PRETTY_NAME=\".+\"")
	reg2 := regexp.MustCompile("\".+\"")
	dstrBytes := reg2.Find(reg1.Find(fData))
	distribution := string(dstrBytes[1 : len(dstrBytes)-1])

	return fmt.Sprintf("%s kernel: %s", distribution, kernel), nil
}
<|endoftext|>


<commit_before>//go:build integration

package cmd

import (
	"fmt"
	"github.com/qiniu/qshell/v2/cmd_test/test"
	"github.com/qiniu/qshell/v2/iqshell/common/account"
	"testing"
)

var accessKey = test.AccessKey
var secretKey = test.SecretKey

func TestUser(t *testing.T) {
	defer func() {
		if err := account.SetAccountToLocalFile(account.Account{
			Name:      "QShell",
			AccessKey: accessKey,
			SecretKey: secretKey,
		}); err != nil {
			fmt.Printf("user set error:%v", err)
		}
	}()

	if test.Local {
		return
	}

	if test.ShouldTestUser {
		TestUserIntegration(t)
	}
}

func TestUserIntegration(t *testing.T) {
	success, err := cleanUser()
	if len(err) > 0 || !success {
		t.Fatal("clean first not success:", err)
	}

	has, err := hasUser()
	if len(err) > 0 || has {
		t.Fatal("shouldn't has user after clean first:", err)
	}

	// 增加 1
	userName := "test_add_1"
	success, err = addUser(userName, accessKey, secretKey)
	if len(err) > 0 || !success {
		t.Fatal("add user not success:", err)
	}

	has, err = hasUser()
	if len(err) > 0 || !has {
		t.Fatal("should has user after add:", err)
	}

	has, err = containUser(userName)
	if len(err) > 0 || !has {
		t.Fatal("should contain user after add:", err)
	}

	has, err = lookupUser(userName)
	if len(err) > 0 || !has {
		t.Fatal("should lookup user success after add:", err)
	}

	is, err := currentUserIs(userName)
	if len(err) > 0 || !is {
		t.Fatal("should change current user after add:", err)
	}

	// 增加 2
	userName = "test_add_2"
	success, err = addUserWithLongOptions(userName, accessKey, secretKey)
	if len(err) > 0 || !success {
		t.Fatal("add user 2 not success:", err)
	}

	has, err = hasUser()
	if len(err) > 0 || !has {
		t.Fatal("should has user after add 2:", err)
	}

	has, err = containUser(userName)
	if len(err) > 0 || !has {
		t.Fatal("shouldn contain user after add 2:", err)
	}

	has, err = lookupUser(userName)
	if len(err) > 0 || !has {
		t.Fatal("should lookup user success after add 2:", err)
	}

	is, err = currentUserIs(userName)
	if len(err) > 0 || !is {
		t.Fatal("should change current user after add 2:", err)
	}

	// 改
	userName = "test_add_1"
	success, err = changeCurrentUser(userName)
	if len(err) > 0 || !success {
		t.Fatal("change current user not success:", err)
	}

	is, err = currentUserIs(userName)
	if len(err) > 0 || !is {
		t.Fatal("should change current user after change :", err)
	}

	// 删除
	success, err = deleteUser(userName)
	if len(err) > 0 || !success {
		t.Fatal("remove user not success:", err)
	}

	has, err = lookupUser(userName)
	if len(err) > 0 || has {
		t.Fatal("shouldn't lookup user success after remove:", err)
	}

	is, err = currentUserIs(userName)
	if len(err) > 0 || !is {
		t.Fatal("shouldn't change current user after delete current user :", err)
	}

	// 清除
	success, err = cleanUser()
	if len(err) > 0 || !success {
		t.Fatal("clean end not success:", err)
	}

	has, err = hasUser()
	if len(err) > 0 || has {
		t.Fatal("shouldn't has user after clean end:", err)
	}

	// 添加测试账号
	userName = "QShell"
	success, err = addUser(userName, accessKey, secretKey)
	if len(err) > 0 || !success {
		t.Fatal("add user not success:", err)
	}
}
<commit_msg>optimize test case<commit_after>//go:build integration

package cmd

import (
	"github.com/qiniu/qshell/v2/cmd"
	"github.com/qiniu/qshell/v2/cmd_test/test"
	"os"
	"testing"
)

var accessKey = test.AccessKey
var secretKey = test.SecretKey

func TestUser(t *testing.T) {
	defer func() {
		os.Args = []string{"qshell", accessKey, secretKey, "QShell"}
		cmd.Execute()
	}()

	if test.Local {
		return
	}

	if test.ShouldTestUser {
		TestUserIntegration(t)
	}
}

func TestUserIntegration(t *testing.T) {
	success, err := cleanUser()
	if len(err) > 0 || !success {
		t.Fatal("clean first not success:", err)
	}

	has, err := hasUser()
	if len(err) > 0 || has {
		t.Fatal("shouldn't has user after clean first:", err)
	}

	// 增加 1
	userName := "test_add_1"
	success, err = addUser(userName, accessKey, secretKey)
	if len(err) > 0 || !success {
		t.Fatal("add user not success:", err)
	}

	has, err = hasUser()
	if len(err) > 0 || !has {
		t.Fatal("should has user after add:", err)
	}

	has, err = containUser(userName)
	if len(err) > 0 || !has {
		t.Fatal("should contain user after add:", err)
	}

	has, err = lookupUser(userName)
	if len(err) > 0 || !has {
		t.Fatal("should lookup user success after add:", err)
	}

	is, err := currentUserIs(userName)
	if len(err) > 0 || !is {
		t.Fatal("should change current user after add:", err)
	}

	// 增加 2
	userName = "test_add_2"
	success, err = addUserWithLongOptions(userName, accessKey, secretKey)
	if len(err) > 0 || !success {
		t.Fatal("add user 2 not success:", err)
	}

	has, err = hasUser()
	if len(err) > 0 || !has {
		t.Fatal("should has user after add 2:", err)
	}

	has, err = containUser(userName)
	if len(err) > 0 || !has {
		t.Fatal("shouldn contain user after add 2:", err)
	}

	has, err = lookupUser(userName)
	if len(err) > 0 || !has {
		t.Fatal("should lookup user success after add 2:", err)
	}

	is, err = currentUserIs(userName)
	if len(err) > 0 || !is {
		t.Fatal("should change current user after add 2:", err)
	}

	// 改
	userName = "test_add_1"
	success, err = changeCurrentUser(userName)
	if len(err) > 0 || !success {
		t.Fatal("change current user not success:", err)
	}

	is, err = currentUserIs(userName)
	if len(err) > 0 || !is {
		t.Fatal("should change current user after change :", err)
	}

	// 删除
	success, err = deleteUser(userName)
	if len(err) > 0 || !success {
		t.Fatal("remove user not success:", err)
	}

	has, err = lookupUser(userName)
	if len(err) > 0 || has {
		t.Fatal("shouldn't lookup user success after remove:", err)
	}

	is, err = currentUserIs(userName)
	if len(err) > 0 || !is {
		t.Fatal("shouldn't change current user after delete current user :", err)
	}

	// 清除
	success, err = cleanUser()
	if len(err) > 0 || !success {
		t.Fatal("clean end not success:", err)
	}

	has, err = hasUser()
	if len(err) > 0 || has {
		t.Fatal("shouldn't has user after clean end:", err)
	}

	// 添加测试账号
	userName = "QShell"
	success, err = addUser(userName, accessKey, secretKey)
	if len(err) > 0 || !success {
		t.Fatal("add user not success:", err)
	}
}
<|endoftext|>


<commit_before>package model

const (
	// Value of undefined id
	UNDEFINED_ID = -1
	// Value of undefined string
	UNDEFINED_STRING = "<UNDEFINED>"
)

// Represents the request for ping task by NQM agent
type NqmPingTaskRequest struct {
	// The connection id of agent(used to identify task configruation)
	ConnectionId string
	// The hostname of agent
	Hostname string
	// The IP address of agent
	// Could be IPv4 or IPv6 format
	IpAddress string
}

// Represents the response for ping task requested from NQM agent
//
// If NeedPing is false, Targets and Command would be empty array
type NqmPingTaskResponse struct {
	// Whether or not the task should be performed
	NeedPing bool

	// The list of target hosts to be probed(ping)
	Targets []NqmTarget

	// The command/arguments of command to be executed
	Command []string
}

// Represents the data of agent
type NqmAgent struct {
	// The id of agent
	Id int

	// The id of ISP, UNDEFINED_ID means there is not such data for this target
	IspId int16
	// The id of province, UNDEFINED_ID means there is not such data for this target
	ProvinceId int16
	// The id of city, UNDEFINED_ID means there is not such data for this target
	CityId int16
}

// Represents the data of target used by NQM agent
type NqmTarget struct {
	// The id of target
	Id int

	// The IP address or FQDN used by ping command
	Host string

	// The id of ISP, UNDEFINED_ID means there is not such data for this target
	IspId int16
	// The id of province, UNDEFINED_ID means there is not such data for this target
	ProvinceId int16
	// The id of city, UNDEFINED_ID means there is not such data for this target
	CityId int16
	// The tag of the target, UNDEFINED_STRING means no such data for this target
	NameTag string
}
<commit_msg>[OWL-277] Add validation and missed property of RPC response<commit_after>package model

import (
	"fmt"
)

const (
	// Value of undefined id
	UNDEFINED_ID = -1

	UNDEFINED_ISP_ID = int16(UNDEFINED_ID)
	UNDEFINED_PROVINCE_ID = int16(UNDEFINED_ID)
	UNDEFINED_CITY_ID = int16(UNDEFINED_ID)

	// Value of undefined string
	UNDEFINED_STRING = "<UNDEFINED>"
)

// Represents the request for ping task by NQM agent
type NqmPingTaskRequest struct {
	// The connection id of agent(used to identify task configruation)
	ConnectionId string `valid:"required"`
	// The hostname of agent
	Hostname string `valid:"required"`
	// The IP address of agent
	// Could be IPv4 or IPv6 format
	IpAddress string `valid:"required"`
}

// Represents the response for ping task requested from NQM agent
//
// If NeedPing is false, Targets and Command would be empty array
type NqmPingTaskResponse struct {
	// Whether or not the task should be performed
	NeedPing bool

	// The data of agent
	// nil if there is no need for ping
	Agent *NqmAgent

	// The list of target hosts to be probed(ping)
	// nil if there is no need for ping
	Targets []NqmTarget

	// The command/arguments of command to be executed
	// nil if there is no need for ping
	Command []string
}

// Represents the data of agent
type NqmAgent struct {
	// The id of agent
	Id int

	// The id of ISP, UNDEFINED_ID means there is not such data for this target
	IspId int16
	// The id of province, UNDEFINED_ID means there is not such data for this target
	ProvinceId int16
	// The id of city, UNDEFINED_ID means there is not such data for this target
	CityId int16
}

// Represents the data of target used by NQM agent
type NqmTarget struct {
	// The id of target
	Id int

	// The IP address or FQDN used by ping command
	Host string

	// The id of ISP, UNDEFINED_ID means there is not such data for this target
	IspId int16
	// The id of province, UNDEFINED_ID means there is not such data for this target
	ProvinceId int16
	// The id of city, UNDEFINED_ID means there is not such data for this target
	CityId int16
	// The tag of the target, UNDEFINED_STRING means no such data for this target
	NameTag string
}

func (target NqmTarget) String() string {
	return fmt.Sprintf(
		"Id: [%d] Host: [%s] IspId: [%d] ProvinceId: [%d], City: [%d], Name tag: [%s]",
		target.Id, target.Host, target.IspId, target.ProvinceId, target.CityId, target.NameTag,
	)
}
<|endoftext|>


<commit_before>package node

import (
	"fmt"
)

type fromChan int

const (
	fromMerge fromChan = iota
	fromControl
	fromInput
	fromStatus
)

const (
	blockIdKey = "block_id"
)

// handleFromMain handles messages received in the main func.
// The primary purpose is to intercept the messages,
// run my list of commands, then potentially wait for
// the commands to finish before forwarding the message.
//
// Note that a side effect is that messages can get
// lost, if I happen to receive a new one while waiting
// to hear the response from my commands after a previous
// message.
type handleFromMain struct {
	owner     Owner
	proc      *process
	status    chan execfini
	needs_run bool
	in_stop   bool
	stop_msg  Msg
	// Solely so I can send a msg down the pipe. Should be a cleaner way.
	ex *Exec
}

func newHandleFromMain(owner Owner, proc *process, status chan execfini, ex *Exec) *handleFromMain {
	return &handleFromMain{owner, proc, status, false, false, Msg{}, ex}
}

func (h *handleFromMain) close() {
	h.proc.close()
}

func (h *handleFromMain) handleMsg(msg *Msg, from fromChan) {
	if msg == nil {
		return
	}
	if from == fromControl {
		h.handleFromControl(msg)
	} else if from == fromInput {
		h.handleFromInput(msg)
	}
}

func (h *handleFromMain) handleFromControl(msg *Msg) {
	cmd := CmdFromMsg(*msg)
	fmt.Println("exec control msg", *msg, "cmd", cmd)
	if cmd != nil {
		if cmd.Method == cmdStop {
			fmt.Println("Got stop for", cmd)
			if !h.proc.isRunning() {
				reply := Cmd{Method: cmdStopReply, TargetId: msg.SenderId}
				rmsg := reply.AsMsg()
				rmsg.SetInt(blockIdKey, msg.MustGetInt(blockIdKey))
				h.owner.SendMsg(rmsg, msg.SenderId)
			} else {
				h.in_stop = true
				h.stop_msg = *msg
				h.proc.stop()
			}
		}
	}
}

func (h *handleFromMain) handleFromInput(msg *Msg) {
	fmt.Println("exec msg", msg)
	if !h.proc.isRunning() {
		h.proc.run(h.status)
	} else {
		h.needs_run = true
	}
}

func (h *handleFromMain) handleFini(fini execfini, from fromChan) {
	if from == fromStatus {
		h.handleFromStatus(fini)
	}
}

func (h *handleFromMain) handleFromStatus(fini execfini) {
	fmt.Println("run fini", fini)
	h.proc.finished(fini)
	if h.in_stop {
		h.in_stop = false
		reply := Cmd{Method: cmdStopReply, TargetId: h.stop_msg.SenderId}
		rmsg := reply.AsMsg()
		rmsg.SetInt(blockIdKey, h.stop_msg.MustGetInt(blockIdKey))
		h.owner.SendMsg(rmsg, h.stop_msg.SenderId)
	} else if h.needs_run {
		h.needs_run = false
		h.proc.run(h.status)
	} else if fini.err == nil {
		// Process completed successfully
		h.ex.SendMsg(Msg{})
	}
}

// handleFromCmds handles messages received in the command func.
// The primary purpose is to intercept the messages,
// run my list of commands, then potentially wait for
// the commands to finish before forwarding the message.
//
// Note that a side effect is that messages can get
// lost, if I happen to receive a new one while waiting
// to hear the response from my commands after a previous
// message.
type handleFromCmds struct {
	// Send commands in blocks, and wait to hear back from all members
	// of a block before proceeding. As soon as we start a new block, the previous is discarded,
	block_id   int
	block_size int
	block_msg  Msg
	// I don't currently have a "message empty" state, and I don't want to store the pointer, so use this
	block_has_msg bool
	owner         Owner
	controlId     Id
	cmdList       []Cmd
	cmds          chan Msg
}

func newHandleFromCmds(owner Owner, controlId Id, cmdList []Cmd, cmds chan Msg) *handleFromCmds {
	return &handleFromCmds{1, 0, Msg{}, false, owner, controlId, cmdList, cmds}
}

func (h *handleFromCmds) handle(msg *Msg, from fromChan) {
	if msg == nil {
		return
	}
	if from == fromMerge {
		h.handleFromMerge(msg)
	} else if from == fromControl {
		h.handleFromControl(msg)
	}
}

func (h *handleFromCmds) handleFromMerge(msg *Msg) {
	h.block_id++
	h.block_size = 0
	h.block_msg = *msg
	h.block_has_msg = true
	for _, v := range h.cmdList {
		m := v.AsMsg()
		m.SenderId = h.controlId
		m.SetInt(blockIdKey, h.block_id)
		err := h.owner.SendMsg(m, v.TargetId)
		fmt.Println("sent stop err", err, "cmd", v, "controlId", h.controlId)
		if err == nil && v.Reply {
			h.block_size++
		}
	}
	// If I'm not waiting to hear back from anyone then just send the message
	if h.block_size <= 0 {
		fmt.Println("Send immediate")
		h.send(msg)
	}
}

func (h *handleFromCmds) handleFromControl(msg *Msg) {
	cmd := CmdFromMsg(*msg)
	if cmd != nil {
		if cmd.Method == cmdStopReply && h.block_id == msg.MustGetInt(blockIdKey) {
			h.block_size--
			if h.block_size == 0 && h.block_has_msg {
				fmt.Println("Send delayed")
				h.send(&h.block_msg)
			}
		}
	}
}

func (h *handleFromCmds) send(msg *Msg) {
	if msg != nil {
		h.cmds <- *msg
		h.block_has_msg = false
	}
}
<commit_msg>Enable hosting mode -- apps that might have crashed are restarted.<commit_after>package node

import (
	"fmt"
)

type fromChan int

const (
	fromMerge fromChan = iota
	fromControl
	fromInput
	fromStatus
)

const (
	blockIdKey = "block_id"
)

// handleFromMain handles messages received in the main func.
// The primary purpose is to intercept the messages,
// run my list of commands, then potentially wait for
// the commands to finish before forwarding the message.
//
// Note that a side effect is that messages can get
// lost, if I happen to receive a new one while waiting
// to hear the response from my commands after a previous
// message.
type handleFromMain struct {
	owner     Owner
	proc      *process
	status    chan execfini
	needs_run bool
	in_stop   bool
	stop_msg  Msg
	// Solely so I can send a msg down the pipe. Should be a cleaner way.
	ex *Exec
}

func newHandleFromMain(owner Owner, proc *process, status chan execfini, ex *Exec) *handleFromMain {
	return &handleFromMain{owner, proc, status, false, false, Msg{}, ex}
}

func (h *handleFromMain) close() {
	h.proc.close()
}

func (h *handleFromMain) handleMsg(msg *Msg, from fromChan) {
	if msg == nil {
		return
	}
	if from == fromControl {
		h.handleFromControl(msg)
	} else if from == fromInput {
		h.handleFromInput(msg)
	}
}

func (h *handleFromMain) handleFromControl(msg *Msg) {
	cmd := CmdFromMsg(*msg)
	fmt.Println("exec control msg", *msg, "cmd", cmd)
	if cmd != nil {
		if cmd.Method == cmdStop {
			fmt.Println("Got stop for", cmd)
			if !h.proc.isRunning() {
				reply := Cmd{Method: cmdStopReply, TargetId: msg.SenderId}
				rmsg := reply.AsMsg()
				rmsg.SetInt(blockIdKey, msg.MustGetInt(blockIdKey))
				h.owner.SendMsg(rmsg, msg.SenderId)
			} else {
				h.in_stop = true
				h.stop_msg = *msg
				h.proc.stop()
			}
		}
	}
}

func (h *handleFromMain) handleFromInput(msg *Msg) {
	fmt.Println("exec msg", msg)
	if !h.proc.isRunning() {
		h.proc.run(h.status)
	} else {
		h.needs_run = true
	}
}

func (h *handleFromMain) handleFini(fini execfini, from fromChan) {
	if from == fromStatus {
		h.handleFromStatus(fini)
	}
}

func (h *handleFromMain) handleFromStatus(fini execfini) {
	fmt.Println("run fini", fini)
	h.proc.finished(fini)
	needs_run := false
	if h.in_stop {
		h.in_stop = false
		reply := Cmd{Method: cmdStopReply, TargetId: h.stop_msg.SenderId}
		rmsg := reply.AsMsg()
		rmsg.SetInt(blockIdKey, h.stop_msg.MustGetInt(blockIdKey))
		h.owner.SendMsg(rmsg, h.stop_msg.SenderId)
	} else if h.needs_run {
		h.needs_run = false
		needs_run = true;
	} else if fini.err == nil {
		// Process completed successfully
		h.ex.SendMsg(Msg{})
		if h.ex.Rerun {
			needs_run = true
		}
	} else {
		// Some sort of error -- possibly a crash.
		if h.ex.Rerun {
			needs_run = true
		}
	}

	if needs_run {
		h.proc.run(h.status)
	}
}

// handleFromCmds handles messages received in the command func.
// The primary purpose is to intercept the messages,
// run my list of commands, then potentially wait for
// the commands to finish before forwarding the message.
//
// Note that a side effect is that messages can get
// lost, if I happen to receive a new one while waiting
// to hear the response from my commands after a previous
// message.
type handleFromCmds struct {
	// Send commands in blocks, and wait to hear back from all members
	// of a block before proceeding. As soon as we start a new block, the previous is discarded,
	block_id   int
	block_size int
	block_msg  Msg
	// I don't currently have a "message empty" state, and I don't want to store the pointer, so use this
	block_has_msg bool
	owner         Owner
	controlId     Id
	cmdList       []Cmd
	cmds          chan Msg
}

func newHandleFromCmds(owner Owner, controlId Id, cmdList []Cmd, cmds chan Msg) *handleFromCmds {
	return &handleFromCmds{1, 0, Msg{}, false, owner, controlId, cmdList, cmds}
}

func (h *handleFromCmds) handle(msg *Msg, from fromChan) {
	if msg == nil {
		return
	}
	if from == fromMerge {
		h.handleFromMerge(msg)
	} else if from == fromControl {
		h.handleFromControl(msg)
	}
}

func (h *handleFromCmds) handleFromMerge(msg *Msg) {
	h.block_id++
	h.block_size = 0
	h.block_msg = *msg
	h.block_has_msg = true
	for _, v := range h.cmdList {
		m := v.AsMsg()
		m.SenderId = h.controlId
		m.SetInt(blockIdKey, h.block_id)
		err := h.owner.SendMsg(m, v.TargetId)
		fmt.Println("sent stop err", err, "cmd", v, "controlId", h.controlId)
		if err == nil && v.Reply {
			h.block_size++
		}
	}
	// If I'm not waiting to hear back from anyone then just send the message
	if h.block_size <= 0 {
		fmt.Println("Send immediate")
		h.send(msg)
	}
}

func (h *handleFromCmds) handleFromControl(msg *Msg) {
	cmd := CmdFromMsg(*msg)
	if cmd != nil {
		if cmd.Method == cmdStopReply && h.block_id == msg.MustGetInt(blockIdKey) {
			h.block_size--
			if h.block_size == 0 && h.block_has_msg {
				fmt.Println("Send delayed")
				h.send(&h.block_msg)
			}
		}
	}
}

func (h *handleFromCmds) send(msg *Msg) {
	if msg != nil {
		h.cmds <- *msg
		h.block_has_msg = false
	}
}
<|endoftext|>


<commit_before>package moka_test

import (
	"github.com/gcapizzi/moka"
	. "github.com/gcapizzi/moka/syntax"

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

var _ = Describe("Moka", func() {
	var collaborator CollaboratorDouble
	var subject Subject

	BeforeEach(func() {
		collaborator = NewCollaboratorDouble()
		subject = NewSubject(collaborator)
	})

	It("allows to stub a method on a double", func() {
		AllowDouble(collaborator).To(ReceiveCallTo("Query").With("arg").AndReturn("result"))
		Expect(subject.DelegateQuery("arg")).To(Equal("result"))
	})

	It("allows to mock a method on a double", func() {
		ExpectDouble(collaborator).To(ReceiveCallTo("Command").With("arg").AndReturn("result"))

		result := subject.DelegateCommand("arg")
		Expect(result).To(Equal("result"))

		VerifyCalls(collaborator)
	})
})

type Collaborator interface {
	Query(string) string
	Command(string) string
}

type CollaboratorDouble struct {
	moka.Double
}

func NewCollaboratorDouble() CollaboratorDouble {
	return CollaboratorDouble{Double: moka.NewStrictDouble()}
}

func (d CollaboratorDouble) Query(arg string) string {
	return d.Call("Query", arg)[0].(string)
}

func (d CollaboratorDouble) Command(arg string) string {
	return d.Call("Command", arg)[0].(string)
}

type Subject struct {
	collaborator Collaborator
}

func NewSubject(collaborator Collaborator) Subject {
	return Subject{collaborator: collaborator}
}

func (s Subject) DelegateQuery(arg string) string {
	return s.collaborator.Query(arg)
}

func (s Subject) DelegateCommand(arg string) string {
	return s.collaborator.Command(arg)
}
<commit_msg>Always VerifyCalls in moka_test<commit_after>package moka_test

import (
	"github.com/gcapizzi/moka"
	. "github.com/gcapizzi/moka/syntax"

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

var _ = Describe("Moka", func() {
	var collaborator CollaboratorDouble
	var subject Subject

	BeforeEach(func() {
		collaborator = NewCollaboratorDouble()
		subject = NewSubject(collaborator)
	})

	AfterEach(func() {
		VerifyCalls(collaborator)
	})

	It("allows to stub a method on a double", func() {
		AllowDouble(collaborator).To(ReceiveCallTo("Query").With("arg").AndReturn("result"))
		Expect(subject.DelegateQuery("arg")).To(Equal("result"))
	})

	It("allows to mock a method on a double", func() {
		ExpectDouble(collaborator).To(ReceiveCallTo("Command").With("arg").AndReturn("result"))

		result := subject.DelegateCommand("arg")
		Expect(result).To(Equal("result"))
	})
})

type Collaborator interface {
	Query(string) string
	Command(string) string
}

type CollaboratorDouble struct {
	moka.Double
}

func NewCollaboratorDouble() CollaboratorDouble {
	return CollaboratorDouble{Double: moka.NewStrictDouble()}
}

func (d CollaboratorDouble) Query(arg string) string {
	return d.Call("Query", arg)[0].(string)
}

func (d CollaboratorDouble) Command(arg string) string {
	return d.Call("Command", arg)[0].(string)
}

type Subject struct {
	collaborator Collaborator
}

func NewSubject(collaborator Collaborator) Subject {
	return Subject{collaborator: collaborator}
}

func (s Subject) DelegateQuery(arg string) string {
	return s.collaborator.Query(arg)
}

func (s Subject) DelegateCommand(arg string) string {
	return s.collaborator.Command(arg)
}
<|endoftext|>


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

// +build !nontp

package collector

import (
	"flag"
	"fmt"
	"time"

	"github.com/beevik/ntp"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
)

var (
	ntpServer          = flag.String("collector.ntp.server", "", "NTP server to use for ntp collector.")
	ntpProtocolVersion = flag.Int("collector.ntp.protocol-version", 4, "NTP protocol version")
)

type ntpCollector struct {
	drift prometheus.Gauge
}

func init() {
	Factories["ntp"] = NewNtpCollector
}

// Takes a prometheus registry and returns a new Collector exposing
// the offset between ntp and the current system time.
func NewNtpCollector() (Collector, error) {
	if *ntpServer == "" {
		return nil, fmt.Errorf("no NTP server specifies, see --ntpServer")
	}
	if *ntpProtocolVersion < 2 || *ntpProtocolVersion > 4 {
		return nil, fmt.Errorf("invalid NTP protocol version %d; must be 2, 3, or 4")
	}

	return &ntpCollector{
		drift: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: Namespace,
			Name:      "ntp_drift_seconds",
			Help:      "Time between system time and ntp time.",
		}),
	}, nil
}

func (c *ntpCollector) Update(ch chan<- prometheus.Metric) (err error) {
	t, err := ntp.TimeV(*ntpServer, byte(*ntpProtocolVersion))
	if err != nil {
		return fmt.Errorf("couldn't get NTP drift: %s", err)
	}
	drift := t.Sub(time.Now())
	log.Debugf("Set ntp_drift_seconds: %f", drift.Seconds())
	c.drift.Set(drift.Seconds())
	c.drift.Collect(ch)
	return err
}
<commit_msg>Add missing argument to fmt.Errorf function in 'ntp' collector<commit_after>// Copyright 2015 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build !nontp

package collector

import (
	"flag"
	"fmt"
	"time"

	"github.com/beevik/ntp"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
)

var (
	ntpServer          = flag.String("collector.ntp.server", "", "NTP server to use for ntp collector.")
	ntpProtocolVersion = flag.Int("collector.ntp.protocol-version", 4, "NTP protocol version")
)

type ntpCollector struct {
	drift prometheus.Gauge
}

func init() {
	Factories["ntp"] = NewNtpCollector
}

// Takes a prometheus registry and returns a new Collector exposing
// the offset between ntp and the current system time.
func NewNtpCollector() (Collector, error) {
	if *ntpServer == "" {
		return nil, fmt.Errorf("no NTP server specifies, see --ntpServer")
	}
	if *ntpProtocolVersion < 2 || *ntpProtocolVersion > 4 {
		return nil, fmt.Errorf("invalid NTP protocol version %d; must be 2, 3, or 4", *ntpProtocolVersion)
	}

	return &ntpCollector{
		drift: prometheus.NewGauge(prometheus.GaugeOpts{
			Namespace: Namespace,
			Name:      "ntp_drift_seconds",
			Help:      "Time between system time and ntp time.",
		}),
	}, nil
}

func (c *ntpCollector) Update(ch chan<- prometheus.Metric) (err error) {
	t, err := ntp.TimeV(*ntpServer, byte(*ntpProtocolVersion))
	if err != nil {
		return fmt.Errorf("couldn't get NTP drift: %s", err)
	}
	drift := t.Sub(time.Now())
	log.Debugf("Set ntp_drift_seconds: %f", drift.Seconds())
	c.drift.Set(drift.Seconds())
	c.drift.Collect(ch)
	return err
}
<|endoftext|>


<commit_before>package message

import (
	"errors"
	"strings"
)

type Message struct {
	Command string
	Options []string
}

var aliases = map[string]string{}

var optionCounts = map[string]int{
	"ADMIN":    1,
	"AWAY":     1,
	"CNOTICE":  3,
	"CPRIVMSG": 3,
	"CONNECT":  3,
	"DIE":      0,
	"ERROR":    1,
	"HELP":     0,
	"INFO":     1,
	"INVITE":   2,
	"ISON":     1,
	"JOIN":     2,
	"KICK":     3,
	"KILL":     2,
	"KNOCK":    2,
	"LINKS":    2,
	"LIST":     2,
	"LUSERS":   2,
	"MODE":     3,
	"MOTD":     1,
	"NAMES":    2,
	"NICK":     2,
	"NOTICE":   2,
	"OPER":     2,
	"PART":     2,
	"PASS":     1,
	"PING":     2,
	"PONG":     2,
	"PRIVMSG":  2,
	"QUIT":     1,
	"REHASH":   0,
	"RESTART":  0,
	"RULES":    0,
	"SERVER":   3,
	"SERVICE":  6,
	"SERVLIST": 2,
	"SQUERY":   2,
	"SQUIT":    2,
	"SETNAME":  1,
	"SILENCE":  1,
	"STATS":    2,
	"SUMMON":   3,
	"TIME":     1,
	"TOPIC":    2,
	"TRACE":    1,
	"USER":     4,
	"USERHOST": 1,
	"USERIP":   1,
	"USERS":    1,
	"VERSION":  1,
	"WALLOPS":  1,
	"WATCH":    1,
	"WHO":      2,
	"WHOIS":    2,
	"WHOWAS":   3,
}

func Marshal(m Message) string {

	fullCmd := []string{}

	//get correct command
	cmd := strings.ToUpper(m.Command)
	if aliases[cmd] != "" {
		cmd = aliases[cmd]
	}

	//check for multi-word last option
	finalOption := ""
	if len(strings.Split(m.Options[len(m.Options)-1], " ")) > 1 {
		finalOption = ":" + m.Options[len(m.Options)-1]
		m.Options = m.Options[:len(m.Options)-1]
	}

	options := strings.Join(m.Options, " ")
	fullCmd = append(fullCmd, cmd)
	if options != "" {
		fullCmd = append(fullCmd, options)
	}
	if finalOption != "" {
		fullCmd = append(fullCmd, finalOption)
	}

	return strings.Join(fullCmd, " ") + "\r\n"
}

func Unmarshal(input string) (Message, error) {

	msg := Message{}

	if input == "" {
		return Message{}, errors.New("Input cannot be empty")
	}

	//remove ending characters
	input = strings.Replace(input, "\n", "", -1)
	input = strings.Replace(input, "\r", "", -1)

	pieces := strings.Split(input, " ")
	if len(pieces) < 2 {
		return Message{}, errors.New("A command is required")
	}

	//Check for prefix
	if string(pieces[0][0]) == ":" {
		//remove extra data
		pieces = pieces[1:]
	}

	//get command
	msg.Command = pieces[0]
	if len(pieces) == 1 {
		return msg, nil
	}
	pieces = pieces[1:]

	//get options
	for pieceIndex, piece := range pieces {
		//Get any multi-word last argument
		if string(piece[0]) == ":" {
			pieces[pieceIndex] = string(piece[1:])
			msg.Options = append(msg.Options, strings.Join(pieces[pieceIndex:], " "))
			break
		}
		msg.Options = append(msg.Options, piece)
	}

	return msg, nil
}

func ParseCommand(input string) (Message, error) {
	msg := Message{}

	if input == "" {
		return Message{}, errors.New("Input cannot be empty")
	}

	pieces := strings.Split(input, " ")

	if pieces[0][0] == '/' {
		msg.Command = strings.ToUpper(pieces[0][1:])
		if aliases[msg.Command] != "" {
			msg.Command = aliases[msg.Command]
		}

		//remove command
		pieces = pieces[1:]
	} else {
		msg.Command = "PRIVMSG"
	}

	numOptions, cmdExists := optionCounts[msg.Command]
	if !cmdExists {
		return Message{}, errors.New("Unknown command")
	}

	for i := 0; i < numOptions; i++ {
		if i > len(pieces)-1 {
			break
		}

		if i == numOptions-1 && numOptions < len(pieces)-1 {
			msg.Options = append(msg.Options, strings.Join(pieces[i:], " "))
			break
		}

		msg.Options = append(msg.Options, pieces[i])
	}

	return msg, nil
}
<commit_msg>Add prefix data to message struct<commit_after>package message

import (
	"errors"
	"strings"
)

type Message struct {
	Prefix  string
	Command string
	Options []string
}

var aliases = map[string]string{}

var optionCounts = map[string]int{
	"ADMIN":    1,
	"AWAY":     1,
	"CNOTICE":  3,
	"CPRIVMSG": 3,
	"CONNECT":  3,
	"DIE":      0,
	"ERROR":    1,
	"HELP":     0,
	"INFO":     1,
	"INVITE":   2,
	"ISON":     1,
	"JOIN":     2,
	"KICK":     3,
	"KILL":     2,
	"KNOCK":    2,
	"LINKS":    2,
	"LIST":     2,
	"LUSERS":   2,
	"MODE":     3,
	"MOTD":     1,
	"NAMES":    2,
	"NICK":     2,
	"NOTICE":   2,
	"OPER":     2,
	"PART":     2,
	"PASS":     1,
	"PING":     2,
	"PONG":     2,
	"PRIVMSG":  2,
	"QUIT":     1,
	"REHASH":   0,
	"RESTART":  0,
	"RULES":    0,
	"SERVER":   3,
	"SERVICE":  6,
	"SERVLIST": 2,
	"SQUERY":   2,
	"SQUIT":    2,
	"SETNAME":  1,
	"SILENCE":  1,
	"STATS":    2,
	"SUMMON":   3,
	"TIME":     1,
	"TOPIC":    2,
	"TRACE":    1,
	"USER":     4,
	"USERHOST": 1,
	"USERIP":   1,
	"USERS":    1,
	"VERSION":  1,
	"WALLOPS":  1,
	"WATCH":    1,
	"WHO":      2,
	"WHOIS":    2,
	"WHOWAS":   3,
}

func Marshal(m Message) string {

	fullCmd := []string{}

	//get correct command
	cmd := strings.ToUpper(m.Command)
	if aliases[cmd] != "" {
		cmd = aliases[cmd]
	}

	//check for multi-word last option
	finalOption := ""
	if len(strings.Split(m.Options[len(m.Options)-1], " ")) > 1 {
		finalOption = ":" + m.Options[len(m.Options)-1]
		m.Options = m.Options[:len(m.Options)-1]
	}

	options := strings.Join(m.Options, " ")

	if m.Prefix != "" {
		fullCmd = append(fullCmd, ":"+m.Prefix)
	}
	fullCmd = append(fullCmd, cmd)
	if options != "" {
		fullCmd = append(fullCmd, options)
	}
	if finalOption != "" {
		fullCmd = append(fullCmd, finalOption)
	}

	return strings.Join(fullCmd, " ") + "\r\n"
}

func Unmarshal(input string) (Message, error) {

	msg := Message{}

	if input == "" {
		return Message{}, errors.New("Input cannot be empty")
	}

	//remove ending characters
	input = strings.Replace(input, "\n", "", -1)
	input = strings.Replace(input, "\r", "", -1)

	pieces := strings.Split(input, " ")
	if len(pieces) < 2 {
		return Message{}, errors.New("A command is required")
	}

	//Check for prefix
	if string(pieces[0][0]) == ":" {
		msg.Prefix = pieces[0]
		//remove prefix data
		pieces = pieces[1:]
	}

	//get command
	msg.Command = pieces[0]
	if len(pieces) == 1 {
		return msg, nil
	}
	pieces = pieces[1:]

	//get options
	for pieceIndex, piece := range pieces {
		//Get any multi-word last argument
		if string(piece[0]) == ":" {
			pieces[pieceIndex] = string(piece[1:])
			msg.Options = append(msg.Options, strings.Join(pieces[pieceIndex:], " "))
			break
		}
		msg.Options = append(msg.Options, piece)
	}

	return msg, nil
}

func ParseCommand(input string) (Message, error) {
	msg := Message{}

	if input == "" {
		return Message{}, errors.New("Input cannot be empty")
	}

	pieces := strings.Split(input, " ")

	if pieces[0][0] == '/' {
		msg.Command = strings.ToUpper(pieces[0][1:])
		if aliases[msg.Command] != "" {
			msg.Command = aliases[msg.Command]
		}

		//remove command
		pieces = pieces[1:]
	} else {
		msg.Command = "PRIVMSG"
	}

	numOptions, cmdExists := optionCounts[msg.Command]
	if !cmdExists {
		return Message{}, errors.New("Unknown command")
	}

	for i := 0; i < numOptions; i++ {
		if i > len(pieces)-1 {
			break
		}

		if i == numOptions-1 && numOptions < len(pieces)-1 {
			msg.Options = append(msg.Options, strings.Join(pieces[i:], " "))
			break
		}

		msg.Options = append(msg.Options, pieces[i])
	}

	return msg, nil
}
<|endoftext|>


<commit_before>package dev

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"

	"gopkg.in/tomb.v2"
)

var ErrUnexpectedExit = errors.New("unexpected exit")

type App struct {
	Name    string
	Port    int
	Command *exec.Cmd

	t tomb.Tomb

	listener net.Listener

	stdout  io.Reader
	lock    sync.Mutex
	pool    *AppPool
	lastUse time.Time
}

func (a *App) Address() string {
	return fmt.Sprintf("localhost:%d", a.Port)
}

func (a *App) watch() error {
	c := make(chan error)

	go func() {
		r := bufio.NewReader(a.stdout)

		for {
			line, err := r.ReadString('\n')
			if line != "" {
				fmt.Fprintf(os.Stdout, "%s[%d]: %s", a.Name, a.Command.Process.Pid, line)
			}

			if err != nil {
				c <- err
				return
			}
		}
	}()

	var err error

	select {
	case err = <-c:
		err = ErrUnexpectedExit
	case <-a.t.Dying():
		a.Command.Process.Kill()
		err = nil
	}

	a.Command.Wait()
	a.pool.remove(a)
	a.listener.Close()

	return err
}

func (a *App) idleMonitor() error {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if a.pool.maybeIdle(a) {
				a.Command.Process.Kill()
			}
			return nil
		case <-a.t.Dying():
			return nil
		}
	}

	return nil
}

func (a *App) UpdateUsed() {
	a.lastUse = time.Now()
}

func LaunchApp(pool *AppPool, name, dir string) (*App, error) {
	// Create a listener socket and inject it
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		return nil, err
	}

	addr := l.Addr().(*net.TCPAddr)

	shell := os.Getenv("SHELL")

	cmd := exec.Command(shell, "-l", "-i", "-c",
		fmt.Sprintf("bundle exec puma -C- --tag puma-dev:%s -b tcp://127.0.0.1:%d",
			name, addr.Port))

	cmd.Dir = dir

	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env,
		fmt.Sprintf("PUMA_INHERIT_0=3:tcp://127.0.0.1:%d", addr.Port))

	tcpListener := l.(*net.TCPListener)
	socket, err := tcpListener.File()
	if err != nil {
		return nil, err
	}

	cmd.ExtraFiles = []*os.File{socket}

	cmd.Stderr = os.Stderr

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	fmt.Printf("! Booted app '%s' on port %d\n", name, addr.Port)

	app := &App{
		Name:     name,
		Port:     addr.Port,
		Command:  cmd,
		listener: l,
		stdout:   stdout,
	}

	app.t.Go(app.watch)
	app.t.Go(app.idleMonitor)

	return app, nil
}

type AppPool struct {
	Dir      string
	IdleTime time.Duration

	lock sync.Mutex
	apps map[string]*App
}

func (a *AppPool) maybeIdle(app *App) bool {
	a.lock.Lock()
	defer a.lock.Unlock()

	diff := time.Since(app.lastUse)
	if diff > a.IdleTime {
		delete(a.apps, app.Name)
		return true
	}

	return false
}

func (a *AppPool) App(name string) (*App, error) {
	a.lock.Lock()
	defer a.lock.Unlock()

	if a.apps == nil {
		a.apps = make(map[string]*App)
	}

	app, ok := a.apps[name]
	if ok {
		app.UpdateUsed()
		return app, nil
	}

	path := filepath.Join(a.Dir, name)

	_, err := os.Stat(path)
	if os.IsNotExist(err) {
		return nil, fmt.Errorf("Unknown app: %s", name)
	}

	app, err = LaunchApp(a, name, path)
	if err != nil {
		return nil, err
	}

	app.pool = a

	app.UpdateUsed()
	a.apps[name] = app

	return app, nil
}

func (a *AppPool) remove(app *App) {
	a.lock.Lock()
	defer a.lock.Unlock()

	fmt.Printf("! Shutdown app '%s'\n", app.Name)

	delete(a.apps, app.Name)
}

func (a *AppPool) Purge() {
	a.lock.Lock()
	defer a.lock.Unlock()

	for _, app := range a.apps {
		app.t.Kill(nil)
	}

	for _, app := range a.apps {
		app.t.Wait()
	}
}
<commit_msg>Execute through the shell to cleanup the processes<commit_after>package dev

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"time"

	"gopkg.in/tomb.v2"
)

var ErrUnexpectedExit = errors.New("unexpected exit")

type App struct {
	Name    string
	Port    int
	Command *exec.Cmd

	t tomb.Tomb

	listener net.Listener

	stdout  io.Reader
	lock    sync.Mutex
	pool    *AppPool
	lastUse time.Time
}

func (a *App) Address() string {
	return fmt.Sprintf("localhost:%d", a.Port)
}

func (a *App) watch() error {
	c := make(chan error)

	go func() {
		r := bufio.NewReader(a.stdout)

		for {
			line, err := r.ReadString('\n')
			if line != "" {
				fmt.Fprintf(os.Stdout, "%s[%d]: %s", a.Name, a.Command.Process.Pid, line)
			}

			if err != nil {
				c <- err
				return
			}
		}
	}()

	var err error

	select {
	case err = <-c:
		err = ErrUnexpectedExit
	case <-a.t.Dying():
		a.Command.Process.Kill()
		err = nil
	}

	a.Command.Wait()
	a.pool.remove(a)
	a.listener.Close()

	return err
}

func (a *App) idleMonitor() error {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			if a.pool.maybeIdle(a) {
				a.Command.Process.Kill()
			}
			return nil
		case <-a.t.Dying():
			return nil
		}
	}

	return nil
}

func (a *App) UpdateUsed() {
	a.lastUse = time.Now()
}

func LaunchApp(pool *AppPool, name, dir string) (*App, error) {
	// Create a listener socket and inject it
	l, err := net.Listen("tcp", ":0")
	if err != nil {
		return nil, err
	}

	addr := l.Addr().(*net.TCPAddr)

	shell := os.Getenv("SHELL")

	cmd := exec.Command(shell, "-l", "-i", "-c",
		fmt.Sprintf("exec bundle exec puma -C- --tag puma-dev:%s -b tcp://127.0.0.1:%d",
			name, addr.Port))

	cmd.Dir = dir

	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env,
		fmt.Sprintf("PUMA_INHERIT_0=3:tcp://127.0.0.1:%d", addr.Port))

	tcpListener := l.(*net.TCPListener)
	socket, err := tcpListener.File()
	if err != nil {
		return nil, err
	}

	cmd.ExtraFiles = []*os.File{socket}

	cmd.Stderr = os.Stderr

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	err = cmd.Start()
	if err != nil {
		return nil, err
	}

	fmt.Printf("! Booted app '%s' on port %d\n", name, addr.Port)

	app := &App{
		Name:     name,
		Port:     addr.Port,
		Command:  cmd,
		listener: l,
		stdout:   stdout,
	}

	app.t.Go(app.watch)
	app.t.Go(app.idleMonitor)

	return app, nil
}

type AppPool struct {
	Dir      string
	IdleTime time.Duration

	lock sync.Mutex
	apps map[string]*App
}

func (a *AppPool) maybeIdle(app *App) bool {
	a.lock.Lock()
	defer a.lock.Unlock()

	diff := time.Since(app.lastUse)
	if diff > a.IdleTime {
		delete(a.apps, app.Name)
		return true
	}

	return false
}

func (a *AppPool) App(name string) (*App, error) {
	a.lock.Lock()
	defer a.lock.Unlock()

	if a.apps == nil {
		a.apps = make(map[string]*App)
	}

	app, ok := a.apps[name]
	if ok {
		app.UpdateUsed()
		return app, nil
	}

	path := filepath.Join(a.Dir, name)

	_, err := os.Stat(path)
	if os.IsNotExist(err) {
		return nil, fmt.Errorf("Unknown app: %s", name)
	}

	app, err = LaunchApp(a, name, path)
	if err != nil {
		return nil, err
	}

	app.pool = a

	app.UpdateUsed()
	a.apps[name] = app

	return app, nil
}

func (a *AppPool) remove(app *App) {
	a.lock.Lock()
	defer a.lock.Unlock()

	fmt.Printf("! Shutdown app '%s'\n", app.Name)

	delete(a.apps, app.Name)
}

func (a *AppPool) Purge() {
	a.lock.Lock()
	defer a.lock.Unlock()

	for _, app := range a.apps {
		app.t.Kill(nil)
	}

	for _, app := range a.apps {
		app.t.Wait()
	}
}
<|endoftext|>