<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"
)

type msg struct {
	Meta             map[string]interface{} `json:"meta,omitempty"`
	Source           string                 `json:"source,omitempty"`
	SystemName       string                 `json:"system-name,omitempty"`
	SubscriptionName string                 `json:"subscription-name,omitempty"`
	Timestamp        int64                  `json:"timestamp,omitempty"`
	Time             *time.Time             `json:"time,omitempty"`
	Prefix           string                 `json:"prefix,omitempty"`
	Updates          []*update              `json:"updates,omitempty"`
	Deletes          []string               `json:"deletes,omitempty"`
}
type update struct {
	Path   string
	Values map[string]interface{} `json:"values,omitempty"`
}

// subscribeCmd represents the subscribe command
var subscribeCmd = &cobra.Command{
	Use:     "subscribe",
	Aliases: []string{"sub"},
	Short:   "subscribe to gnmi updates on targets",

	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()
		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"),
		}

		coll := collector.NewCollector(ctx, 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()
				if err = coll.Subscribe(tn); err != nil {
					if errors.Is(err, context.DeadlineExceeded) {
						logger.Printf("failed to initialize target '%s' timeout (%s) reached", tn, targetsConfig[tn].Timeout)
						return
					}
					logger.Printf("failed to initialize target '%s': %v", tn, err)
				}
			}(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{}{}
			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 := collector.FormatMsg(nil, response, true, "  ")
						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()
		return nil
	},
}

func init() {
	rootCmd.AddCommand(subscribeCmd)
	subscribeCmd.SilenceUsage = true
	subscribeCmd.Flags().StringP("prefix", "", "", "subscribe request prefix")
	subscribeCmd.Flags().StringSliceP("path", "", []string{""}, "subscribe request paths")
	//subscribeCmd.MarkFlagRequired("path")
	subscribeCmd.Flags().Int32P("qos", "q", 20, "qos marking")
	subscribeCmd.Flags().BoolP("updates-only", "", false, "only updates to current state should be sent")
	subscribeCmd.Flags().StringP("mode", "", "stream", "one of: once, stream, poll")
	subscribeCmd.Flags().StringP("stream-mode", "", "target-defined", "one of: on-change, sample, target-defined")
	subscribeCmd.Flags().DurationP("sample-interval", "i", 0,
		"sample interval as a decimal number and a suffix unit, such as \"10s\" or \"1m30s\"")
	subscribeCmd.Flags().BoolP("suppress-redundant", "", false, "suppress redundant update if the subscribed value didn't not change")
	subscribeCmd.Flags().DurationP("heartbeat-interval", "", 0, "heartbeat interval in case suppress-redundant is enabled")
	subscribeCmd.Flags().StringSliceP("model", "", []string{""}, "subscribe request used model(s)")
	subscribeCmd.Flags().Bool("quiet", false, "suppress stdout printing")
	subscribeCmd.Flags().StringP("target", "", "", "subscribe request target")
	//
	viper.BindPFlag("subscribe-prefix", subscribeCmd.LocalFlags().Lookup("prefix"))
	viper.BindPFlag("subscribe-path", subscribeCmd.LocalFlags().Lookup("path"))
	viper.BindPFlag("subscribe-qos", subscribeCmd.LocalFlags().Lookup("qos"))
	viper.BindPFlag("subscribe-updates-only", subscribeCmd.LocalFlags().Lookup("updates-only"))
	viper.BindPFlag("subscribe-mode", subscribeCmd.LocalFlags().Lookup("mode"))
	viper.BindPFlag("subscribe-stream-mode", subscribeCmd.LocalFlags().Lookup("stream-mode"))
	viper.BindPFlag("subscribe-sample-interval", subscribeCmd.LocalFlags().Lookup("sample-interval"))
	viper.BindPFlag("subscribe-suppress-redundant", subscribeCmd.LocalFlags().Lookup("suppress-redundant"))
	viper.BindPFlag("subscribe-heartbeat-interval", subscribeCmd.LocalFlags().Lookup("heartbeat-interval"))
	viper.BindPFlag("subscribe-sub-model", subscribeCmd.LocalFlags().Lookup("model"))
	viper.BindPFlag("subscribe-quiet", subscribeCmd.LocalFlags().Lookup("quiet"))
	viper.BindPFlag("subscribe-target", subscribeCmd.LocalFlags().Lookup("target"))
}

func getOutputs() (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()
							err := o.Init(ou, logger)
							if err != nil {
								return nil, err
							}
							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)
	paths := viper.GetStringSlice("subscribe-path")
	hi := viper.GetDuration("subscribe-heartbeat-interval")
	si := viper.GetDuration("subscribe-sample-interval")
	qos := viper.GetUint32("subscribe-qos")
	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
	}
	return subscriptions, nil
}
<commit_msg>use MarshalOptions in subscribe cmd (poll)<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"
)

type msg struct {
	Meta             map[string]interface{} `json:"meta,omitempty"`
	Source           string                 `json:"source,omitempty"`
	SystemName       string                 `json:"system-name,omitempty"`
	SubscriptionName string                 `json:"subscription-name,omitempty"`
	Timestamp        int64                  `json:"timestamp,omitempty"`
	Time             *time.Time             `json:"time,omitempty"`
	Prefix           string                 `json:"prefix,omitempty"`
	Updates          []*update              `json:"updates,omitempty"`
	Deletes          []string               `json:"deletes,omitempty"`
}
type update struct {
	Path   string
	Values map[string]interface{} `json:"values,omitempty"`
}

// subscribeCmd represents the subscribe command
var subscribeCmd = &cobra.Command{
	Use:     "subscribe",
	Aliases: []string{"sub"},
	Short:   "subscribe to gnmi updates on targets",

	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()
		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"),
		}

		coll := collector.NewCollector(ctx, 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()
				if err = coll.Subscribe(tn); err != nil {
					if errors.Is(err, context.DeadlineExceeded) {
						logger.Printf("failed to initialize target '%s' timeout (%s) reached", tn, targetsConfig[tn].Timeout)
						return
					}
					logger.Printf("failed to initialize target '%s': %v", tn, err)
				}
			}(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()
		return nil
	},
}

func init() {
	rootCmd.AddCommand(subscribeCmd)
	subscribeCmd.SilenceUsage = true
	subscribeCmd.Flags().StringP("prefix", "", "", "subscribe request prefix")
	subscribeCmd.Flags().StringSliceP("path", "", []string{""}, "subscribe request paths")
	//subscribeCmd.MarkFlagRequired("path")
	subscribeCmd.Flags().Int32P("qos", "q", 20, "qos marking")
	subscribeCmd.Flags().BoolP("updates-only", "", false, "only updates to current state should be sent")
	subscribeCmd.Flags().StringP("mode", "", "stream", "one of: once, stream, poll")
	subscribeCmd.Flags().StringP("stream-mode", "", "target-defined", "one of: on-change, sample, target-defined")
	subscribeCmd.Flags().DurationP("sample-interval", "i", 0,
		"sample interval as a decimal number and a suffix unit, such as \"10s\" or \"1m30s\"")
	subscribeCmd.Flags().BoolP("suppress-redundant", "", false, "suppress redundant update if the subscribed value didn't not change")
	subscribeCmd.Flags().DurationP("heartbeat-interval", "", 0, "heartbeat interval in case suppress-redundant is enabled")
	subscribeCmd.Flags().StringSliceP("model", "", []string{""}, "subscribe request used model(s)")
	subscribeCmd.Flags().Bool("quiet", false, "suppress stdout printing")
	subscribeCmd.Flags().StringP("target", "", "", "subscribe request target")
	//
	viper.BindPFlag("subscribe-prefix", subscribeCmd.LocalFlags().Lookup("prefix"))
	viper.BindPFlag("subscribe-path", subscribeCmd.LocalFlags().Lookup("path"))
	viper.BindPFlag("subscribe-qos", subscribeCmd.LocalFlags().Lookup("qos"))
	viper.BindPFlag("subscribe-updates-only", subscribeCmd.LocalFlags().Lookup("updates-only"))
	viper.BindPFlag("subscribe-mode", subscribeCmd.LocalFlags().Lookup("mode"))
	viper.BindPFlag("subscribe-stream-mode", subscribeCmd.LocalFlags().Lookup("stream-mode"))
	viper.BindPFlag("subscribe-sample-interval", subscribeCmd.LocalFlags().Lookup("sample-interval"))
	viper.BindPFlag("subscribe-suppress-redundant", subscribeCmd.LocalFlags().Lookup("suppress-redundant"))
	viper.BindPFlag("subscribe-heartbeat-interval", subscribeCmd.LocalFlags().Lookup("heartbeat-interval"))
	viper.BindPFlag("subscribe-sub-model", subscribeCmd.LocalFlags().Lookup("model"))
	viper.BindPFlag("subscribe-quiet", subscribeCmd.LocalFlags().Lookup("quiet"))
	viper.BindPFlag("subscribe-target", subscribeCmd.LocalFlags().Lookup("target"))
}

func getOutputs() (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()
							err := o.Init(ou, logger)
							if err != nil {
								return nil, err
							}
							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)
	paths := viper.GetStringSlice("subscribe-path")
	hi := viper.GetDuration("subscribe-heartbeat-interval")
	si := viper.GetDuration("subscribe-sample-interval")
	qos := viper.GetUint32("subscribe-qos")
	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
	}
	return subscriptions, nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"github.com/datawire/consul-x/pkg/consulwatch"
	"github.com/datawire/teleproxy/pkg/k8s"
	"github.com/datawire/teleproxy/pkg/supervisor"
	"github.com/datawire/teleproxy/pkg/watt"
	"github.com/spf13/cobra"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
)

var kubernetesNamespace string
var initialSources = make([]string, 0)
var notifyReceivers = make([]string, 0)
var port int

var rootCmd = &cobra.Command{
	Use:              "watt",
	Short:            "watt",
	Long:             "watt - watch all the things",
	PersistentPreRun: func(cmd *cobra.Command, args []string) {},
	Run:              runWatt,
}

type assembler struct {
	aggregatorNotifyCh <-chan struct{}
	snapshotRequestCh  <-chan getSnapshotRequest
	snapshots          map[int]string
}

func (a *assembler) Work(p *supervisor.Process) error {
	snapshots := make(map[int]string)
	snapshots[1] = `{"TODO": "Rafi!"}`

	p.Ready()
	for {
		select {
		case snapshotRequest := <-a.snapshotRequestCh:
			snapshotRequest.result <- getSnapshotResult{found: true, snapshot: snapshots[1]}
		case <-p.Shutdown():
			p.Logf("shutdown initiated")
			return nil
		}
	}
}

type bootstrappah struct {
	kubernetesResourcesCh <-chan []k8s.Resource
	consulEndpointsCh     <-chan consulwatch.Endpoints
	aggregator            *aggregator
}

func (b *bootstrappah) Work(p *supervisor.Process) error {
	p.Ready()
	requiredConsulServices := make(map[string]*consulwatch.Endpoints)

	var bootstrapped = false
	for {
		select {
		case resources := <-b.kubernetesResourcesCh:
			b.aggregator.setKubernetesResources(resources)
			for _, v := range b.aggregator.kubernetesResources["ConsulResolver"] {
				// this is all kinds of type unsafe most likely
				requiredConsulServices[v.Data()["service"].(string)] = nil
			}

			p.Logf("discovered %d consul resolver configurations", len(requiredConsulServices))
		case endpoints := <-b.consulEndpointsCh:
			requiredConsulServices[endpoints.Service] = &endpoints

			if !MapHasNilValues(requiredConsulServices) {
				bootstrapped = true
			}
		}

		p.Logf("bootstrapped!")
		if bootstrapped {
			break
		}
	}

	for _, v := range requiredConsulServices {
		b.aggregator.updateConsulEndpoints(*v)
	}

	p.Supervisor().Supervise(&supervisor.Worker{
		Name: "aggregator",
		Work: b.aggregator.Work,
	})

	return nil
}

func MapHasNilValues(m map[string]*consulwatch.Endpoints) bool {
	for _, v := range m {
		if v == nil {
			return true
		}
	}

	return false
}

type aggregator struct {
	kubernetesResourcesCh <-chan []k8s.Resource
	kubernetesResources   map[string][]k8s.Resource
	consulEndpointsCh     <-chan consulwatch.Endpoints
	consulEndpoints       map[string]consulwatch.Endpoints
	notifyAssemblerCh     chan<- struct{}
}

func (a *aggregator) Work(p *supervisor.Process) error {
	p.Ready()

	for {
		a.notifyAssemblerCh <- struct{}{}

		select {
		case resources := <-a.kubernetesResourcesCh:
			a.setKubernetesResources(resources)
		case endpoints := <-a.consulEndpointsCh:
			a.updateConsulEndpoints(endpoints)
		}
	}
}

func (a *aggregator) updateConsulEndpoints(endpoints consulwatch.Endpoints) {
	fmt.Println(endpoints)
	a.consulEndpoints[endpoints.Service] = endpoints
}

func (a *aggregator) setKubernetesResources(resources []k8s.Resource) {
	replacement := make(map[string][]k8s.Resource)
	for _, r := range resources {
		kind := r.Kind()

		if isConsulResolver(r) {
			kind = "ConsulResolver" // fake it till you make it baby; this will make lookups quicker
		}

		if _, exists := replacement[kind]; !exists {
			replacement[kind] = make([]k8s.Resource, 0)
		}

		replacement[kind] = append(replacement[kind], r)
	}

	a.kubernetesResources = replacement
}

type consulwatchman struct {
	kubernetes                <-chan []k8s.Resource
	consulEndpointsAggregator chan<- consulwatch.Endpoints
	watched                   map[string]*supervisor.Worker
	ready                     bool
}

func (w *consulwatchman) Work(p *supervisor.Process) error {
	p.Ready()

	for {
		select {
		case resources := <-w.kubernetes:
			found := make(map[string]*supervisor.Worker)
			p.Logf("processing %d kubernetes resources", len(resources))
			for _, r := range resources {
				if isConsulResolver(r) {
					worker, err := w.makeConsulWatcher(r)
					if err != nil {
						p.Logf("failed to create consul watch %v", err)
						continue
					}

					if _, exists := w.watched[worker.Name]; !exists {
						p.Logf("add consul watcher %s\n", worker.Name)
						p.Supervisor().Supervise(worker)
						w.watched[worker.Name] = worker
					}

					found[worker.Name] = worker
				}
			}

			// purge the watches that no longer are needed because they did not come through the in the latest
			// report
			for k, worker := range w.watched {
				if _, exists := found[k]; !exists {
					p.Logf("remove consul watcher %s\n", k)
					worker.Shutdown()
				}
			}

			w.watched = found
		case <-p.Shutdown():
			p.Logf("shutdown initiated")
			return nil
		}
	}
}

func (w *consulwatchman) makeConsulWatcher(r k8s.Resource) (*supervisor.Worker, error) {
	data := r.Data()
	cwm := &watt.ConsulServiceNodeWatchMaker{
		Service:     data["service"].(string),
		Datacenter:  data["datacenter"].(string),
		OnlyHealthy: true,
	}

	cwmFunc, err := cwm.Make(w.consulEndpointsAggregator)
	if err != nil {
		return nil, err
	}

	return &supervisor.Worker{
		Name:  cwm.ID(),
		Work:  cwmFunc,
		Retry: false,
	}, nil
}

type kubewatchman struct {
	namespace string
	kinds     []string
	notify    []chan<- []k8s.Resource
}

func fmtNamespace(ns string) string {
	if ns == "" {
		return "*"
	}

	return ns
}

func (w *kubewatchman) Work(p *supervisor.Process) error {
	kubeAPIWatcher := k8s.NewClient(nil).Watcher()

	for _, kind := range w.kinds {
		p.Logf("adding kubernetes watch for %q in namespace %q", kind, fmtNamespace(kubernetesNamespace))

		watcherFunc := func(ns, kind string) func(watcher *k8s.Watcher) {
			return func(watcher *k8s.Watcher) {
				resources := watcher.List(kind)
				p.Logf("found %d %q in namespace %q", len(resources), kind, fmtNamespace(ns))
				for _, n := range w.notify {
					n <- resources
				}
				p.Logf("sent %q to %d receivers", kind, len(w.notify))
			}
		}

		err := kubeAPIWatcher.WatchNamespace(w.namespace, kind, watcherFunc(w.namespace, kind))

		if err != nil {
			return err
		}
	}

	kubeAPIWatcher.Start()
	p.Ready()

	for {
		select {
		case <-p.Shutdown():
			p.Logf("shutdown initiated")
			kubeAPIWatcher.Stop()
			return nil
		}
	}
}

func isConsulResolver(r k8s.Resource) bool {
	kind := strings.ToLower(r.Kind())
	switch kind {
	case "configmap":
		a := r.Metadata().Annotations()
		if _, ok := a["getambassador.io/consul-resolver"]; ok {
			return true
		}
	}

	return false
}

type apiServer struct {
	port        int
	assemblerCh chan<- getSnapshotRequest
}

type getSnapshotResult struct {
	found    bool
	snapshot string
}

type getSnapshotRequest struct {
	id     int
	result chan<- getSnapshotResult
}

func (s *apiServer) Work(p *supervisor.Process) error {
	p.Ready()

	http.HandleFunc("/snapshots/", func(w http.ResponseWriter, r *http.Request) {
		id, err := strconv.Atoi(strings.TrimPrefix(r.URL.Path, "/snapshots/"))
		if err != nil {
			p.Logf("ID is not an integer")
		}

		resultCh := make(chan getSnapshotResult)
		s.assemblerCh <- getSnapshotRequest{id: id, result: resultCh}

		result := <-resultCh
		if !result.found {
			w.WriteHeader(http.StatusNotFound)
			return
		}

		w.Header().Set("content-type", "application/json")
		if _, err := w.Write([]byte(result.snapshot)); err != nil {
			p.Logf("write snapshot error: %v", err)
		}
	})

	listenHostAndPort := fmt.Sprintf(":%d", port)
	p.Logf("snapshot server listening on: %s", listenHostAndPort)
	return http.ListenAndServe(listenHostAndPort, nil)
}

func init() {
	rootCmd.Flags().StringVarP(&kubernetesNamespace, "namespace", "n", "", "namespace to watch (default: all)")
	rootCmd.Flags().StringSliceVarP(&initialSources, "source", "s", []string{}, "configure an initial static source")
	rootCmd.Flags().StringSliceVar(&notifyReceivers, "notify", []string{}, "invoke the program with the given arguments as a receiver")
	rootCmd.Flags().IntVarP(&port, "port", "p", 7000, "configure the snapshot server port")
}

func runWatt(cmd *cobra.Command, args []string) {
	log.Printf("starting watt...")

	consulWatchmanCh := make(chan []k8s.Resource)
	kubernetesResourceAggregatorCh := make(chan []k8s.Resource)
	consulEndpointsAggregatorCh := make(chan consulwatch.Endpoints)

	notifyAssembler := make(chan struct{})

	// bootstrapper waits for steady state then launches the aggregator
	bootstrappah := &bootstrappah{
		kubernetesResourcesCh: kubernetesResourceAggregatorCh,
		consulEndpointsCh:     consulEndpointsAggregatorCh,
		aggregator: &aggregator{
			kubernetesResourcesCh: kubernetesResourceAggregatorCh,
			kubernetesResources:   make(map[string][]k8s.Resource),
			consulEndpointsCh:     consulEndpointsAggregatorCh,
			consulEndpoints:       make(map[string]consulwatch.Endpoints),
			notifyAssemblerCh:     notifyAssembler,
		},
	}

	kubewatchman := kubewatchman{
		namespace: kubernetesNamespace,
		kinds:     initialSources,
		notify:    []chan<- []k8s.Resource{kubernetesResourceAggregatorCh, consulWatchmanCh},
	}

	consulwatchman := consulwatchman{
		kubernetes:                consulWatchmanCh,
		consulEndpointsAggregator: consulEndpointsAggregatorCh,
		watched:                   make(map[string]*supervisor.Worker),
	}

	snapshotRequestCh := make(chan getSnapshotRequest)
	assembler := &assembler{
		aggregatorNotifyCh: notifyAssembler,
		snapshotRequestCh:  snapshotRequestCh,
	}

	apiServer := &apiServer{
		port:        port,
		assemblerCh: snapshotRequestCh,
	}

	ctx := context.Background()
	s := supervisor.WithContext(ctx)

	s.Supervise(&supervisor.Worker{
		Name:     "kubewatchman",
		Work:     kubewatchman.Work,
		Requires: []string{"bootstrappah"},
	})

	s.Supervise(&supervisor.Worker{
		Name:     "consulwatchman",
		Work:     consulwatchman.Work,
		Requires: []string{"bootstrappah"},
	})

	s.Supervise(&supervisor.Worker{
		Name: "bootstrappah",
		Work: bootstrappah.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name:     "assembler",
		Work:     assembler.Work,
		Requires: []string{"aggregator"},
	})

	s.Supervise(&supervisor.Worker{
		Name:     "api",
		Work:     apiServer.Work,
		Requires: []string{"assembler"},
	})

	if errs := s.Run(); len(errs) > 0 {
		for _, err := range errs {
			log.Println(err)
		}
		os.Exit(1)
	}
}

func main() {
	if err := rootCmd.Execute(); err != nil {
		log.Fatalln(err)
	}
}
<commit_msg>renamted assembler -> invoker<commit_after>package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/datawire/consul-x/pkg/consulwatch"
	"github.com/datawire/teleproxy/pkg/k8s"
	"github.com/datawire/teleproxy/pkg/supervisor"
	"github.com/datawire/teleproxy/pkg/watt"
	"github.com/spf13/cobra"
)

var kubernetesNamespace string
var initialSources = make([]string, 0)
var notifyReceivers = make([]string, 0)
var port int

var rootCmd = &cobra.Command{
	Use:              "watt",
	Short:            "watt",
	Long:             "watt - watch all the things",
	PersistentPreRun: func(cmd *cobra.Command, args []string) {},
	Run:              runWatt,
}

type invoker struct {
	aggregatorNotifyCh <-chan struct{}
	snapshotRequestCh  <-chan getSnapshotRequest
	snapshots          map[int]string
}

func (a *invoker) Work(p *supervisor.Process) error {
	snapshots := make(map[int]string)
	snapshots[1] = `{"TODO": "Rafi!"}`

	p.Ready()
	for {
		select {
		case snapshotRequest := <-a.snapshotRequestCh:
			snapshotRequest.result <- getSnapshotResult{found: true, snapshot: snapshots[1]}
		case <-p.Shutdown():
			p.Logf("shutdown initiated")
			return nil
		}
	}
}

type bootstrappah struct {
	kubernetesResourcesCh <-chan []k8s.Resource
	consulEndpointsCh     <-chan consulwatch.Endpoints
	aggregator            *aggregator
}

func (b *bootstrappah) Work(p *supervisor.Process) error {
	p.Ready()
	requiredConsulServices := make(map[string]*consulwatch.Endpoints)

	var bootstrapped = false
	for {
		select {
		case resources := <-b.kubernetesResourcesCh:
			b.aggregator.setKubernetesResources(resources)
			for _, v := range b.aggregator.kubernetesResources["ConsulResolver"] {
				// this is all kinds of type unsafe most likely
				requiredConsulServices[v.Data()["service"].(string)] = nil
			}

			p.Logf("discovered %d consul resolver configurations", len(requiredConsulServices))
		case endpoints := <-b.consulEndpointsCh:
			requiredConsulServices[endpoints.Service] = &endpoints

			if !MapHasNilValues(requiredConsulServices) {
				bootstrapped = true
			}
		}

		p.Logf("bootstrapped!")
		if bootstrapped {
			break
		}
	}

	for _, v := range requiredConsulServices {
		b.aggregator.updateConsulEndpoints(*v)
	}

	p.Supervisor().Supervise(&supervisor.Worker{
		Name: "aggregator",
		Work: b.aggregator.Work,
	})

	return nil
}

func MapHasNilValues(m map[string]*consulwatch.Endpoints) bool {
	for _, v := range m {
		if v == nil {
			return true
		}
	}

	return false
}

type aggregator struct {
	kubernetesResourcesCh <-chan []k8s.Resource
	kubernetesResources   map[string][]k8s.Resource
	consulEndpointsCh     <-chan consulwatch.Endpoints
	consulEndpoints       map[string]consulwatch.Endpoints
	notifyInvokerCh       chan<- struct{}
}

func (a *aggregator) Work(p *supervisor.Process) error {
	p.Ready()

	for {
		a.notifyInvokerCh <- struct{}{}

		select {
		case resources := <-a.kubernetesResourcesCh:
			a.setKubernetesResources(resources)
		case endpoints := <-a.consulEndpointsCh:
			a.updateConsulEndpoints(endpoints)
		}
	}
}

func (a *aggregator) updateConsulEndpoints(endpoints consulwatch.Endpoints) {
	fmt.Println(endpoints)
	a.consulEndpoints[endpoints.Service] = endpoints
}

func (a *aggregator) setKubernetesResources(resources []k8s.Resource) {
	replacement := make(map[string][]k8s.Resource)
	for _, r := range resources {
		kind := r.Kind()

		if isConsulResolver(r) {
			kind = "ConsulResolver" // fake it till you make it baby; this will make lookups quicker
		}

		if _, exists := replacement[kind]; !exists {
			replacement[kind] = make([]k8s.Resource, 0)
		}

		replacement[kind] = append(replacement[kind], r)
	}

	a.kubernetesResources = replacement
}

type consulwatchman struct {
	kubernetes                <-chan []k8s.Resource
	consulEndpointsAggregator chan<- consulwatch.Endpoints
	watched                   map[string]*supervisor.Worker
	ready                     bool
}

func (w *consulwatchman) Work(p *supervisor.Process) error {
	p.Ready()

	for {
		select {
		case resources := <-w.kubernetes:
			found := make(map[string]*supervisor.Worker)
			p.Logf("processing %d kubernetes resources", len(resources))
			for _, r := range resources {
				if isConsulResolver(r) {
					worker, err := w.makeConsulWatcher(r)
					if err != nil {
						p.Logf("failed to create consul watch %v", err)
						continue
					}

					if _, exists := w.watched[worker.Name]; !exists {
						p.Logf("add consul watcher %s\n", worker.Name)
						p.Supervisor().Supervise(worker)
						w.watched[worker.Name] = worker
					}

					found[worker.Name] = worker
				}
			}

			// purge the watches that no longer are needed because they did not come through the in the latest
			// report
			for k, worker := range w.watched {
				if _, exists := found[k]; !exists {
					p.Logf("remove consul watcher %s\n", k)
					worker.Shutdown()
				}
			}

			w.watched = found
		case <-p.Shutdown():
			p.Logf("shutdown initiated")
			return nil
		}
	}
}

func (w *consulwatchman) makeConsulWatcher(r k8s.Resource) (*supervisor.Worker, error) {
	data := r.Data()
	cwm := &watt.ConsulServiceNodeWatchMaker{
		Service:     data["service"].(string),
		Datacenter:  data["datacenter"].(string),
		OnlyHealthy: true,
	}

	cwmFunc, err := cwm.Make(w.consulEndpointsAggregator)
	if err != nil {
		return nil, err
	}

	return &supervisor.Worker{
		Name:  cwm.ID(),
		Work:  cwmFunc,
		Retry: false,
	}, nil
}

type kubewatchman struct {
	namespace string
	kinds     []string
	notify    []chan<- []k8s.Resource
}

func fmtNamespace(ns string) string {
	if ns == "" {
		return "*"
	}

	return ns
}

func (w *kubewatchman) Work(p *supervisor.Process) error {
	kubeAPIWatcher := k8s.NewClient(nil).Watcher()

	for _, kind := range w.kinds {
		p.Logf("adding kubernetes watch for %q in namespace %q", kind, fmtNamespace(kubernetesNamespace))

		watcherFunc := func(ns, kind string) func(watcher *k8s.Watcher) {
			return func(watcher *k8s.Watcher) {
				resources := watcher.List(kind)
				p.Logf("found %d %q in namespace %q", len(resources), kind, fmtNamespace(ns))
				for _, n := range w.notify {
					n <- resources
				}
				p.Logf("sent %q to %d receivers", kind, len(w.notify))
			}
		}

		err := kubeAPIWatcher.WatchNamespace(w.namespace, kind, watcherFunc(w.namespace, kind))

		if err != nil {
			return err
		}
	}

	kubeAPIWatcher.Start()
	p.Ready()

	for {
		select {
		case <-p.Shutdown():
			p.Logf("shutdown initiated")
			kubeAPIWatcher.Stop()
			return nil
		}
	}
}

func isConsulResolver(r k8s.Resource) bool {
	kind := strings.ToLower(r.Kind())
	switch kind {
	case "configmap":
		a := r.Metadata().Annotations()
		if _, ok := a["getambassador.io/consul-resolver"]; ok {
			return true
		}
	}

	return false
}

type apiServer struct {
	port      int
	invokerCh chan<- getSnapshotRequest
}

type getSnapshotResult struct {
	found    bool
	snapshot string
}

type getSnapshotRequest struct {
	id     int
	result chan<- getSnapshotResult
}

func (s *apiServer) Work(p *supervisor.Process) error {
	p.Ready()

	http.HandleFunc("/snapshots/", func(w http.ResponseWriter, r *http.Request) {
		id, err := strconv.Atoi(strings.TrimPrefix(r.URL.Path, "/snapshots/"))
		if err != nil {
			p.Logf("ID is not an integer")
		}

		resultCh := make(chan getSnapshotResult)
		s.invokerCh <- getSnapshotRequest{id: id, result: resultCh}

		result := <-resultCh
		if !result.found {
			w.WriteHeader(http.StatusNotFound)
			return
		}

		w.Header().Set("content-type", "application/json")
		if _, err := w.Write([]byte(result.snapshot)); err != nil {
			p.Logf("write snapshot error: %v", err)
		}
	})

	listenHostAndPort := fmt.Sprintf(":%d", port)
	p.Logf("snapshot server listening on: %s", listenHostAndPort)
	return http.ListenAndServe(listenHostAndPort, nil)
}

func init() {
	rootCmd.Flags().StringVarP(&kubernetesNamespace, "namespace", "n", "", "namespace to watch (default: all)")
	rootCmd.Flags().StringSliceVarP(&initialSources, "source", "s", []string{}, "configure an initial static source")
	rootCmd.Flags().StringSliceVar(&notifyReceivers, "notify", []string{}, "invoke the program with the given arguments as a receiver")
	rootCmd.Flags().IntVarP(&port, "port", "p", 7000, "configure the snapshot server port")
}

func runWatt(cmd *cobra.Command, args []string) {
	log.Printf("starting watt...")

	consulWatchmanCh := make(chan []k8s.Resource)
	kubernetesResourceAggregatorCh := make(chan []k8s.Resource)
	consulEndpointsAggregatorCh := make(chan consulwatch.Endpoints)

	notifyInvoker := make(chan struct{})

	// bootstrapper waits for steady state then launches the aggregator
	bootstrappah := &bootstrappah{
		kubernetesResourcesCh: kubernetesResourceAggregatorCh,
		consulEndpointsCh:     consulEndpointsAggregatorCh,
		aggregator: &aggregator{
			kubernetesResourcesCh: kubernetesResourceAggregatorCh,
			kubernetesResources:   make(map[string][]k8s.Resource),
			consulEndpointsCh:     consulEndpointsAggregatorCh,
			consulEndpoints:       make(map[string]consulwatch.Endpoints),
			notifyInvokerCh:       notifyInvoker,
		},
	}

	kubewatchman := kubewatchman{
		namespace: kubernetesNamespace,
		kinds:     initialSources,
		notify:    []chan<- []k8s.Resource{kubernetesResourceAggregatorCh, consulWatchmanCh},
	}

	consulwatchman := consulwatchman{
		kubernetes:                consulWatchmanCh,
		consulEndpointsAggregator: consulEndpointsAggregatorCh,
		watched:                   make(map[string]*supervisor.Worker),
	}

	snapshotRequestCh := make(chan getSnapshotRequest)
	invoker := &invoker{
		aggregatorNotifyCh: notifyInvoker,
		snapshotRequestCh:  snapshotRequestCh,
	}

	apiServer := &apiServer{
		port:      port,
		invokerCh: snapshotRequestCh,
	}

	ctx := context.Background()
	s := supervisor.WithContext(ctx)

	s.Supervise(&supervisor.Worker{
		Name:     "kubewatchman",
		Work:     kubewatchman.Work,
		Requires: []string{"bootstrappah"},
	})

	s.Supervise(&supervisor.Worker{
		Name:     "consulwatchman",
		Work:     consulwatchman.Work,
		Requires: []string{"bootstrappah"},
	})

	s.Supervise(&supervisor.Worker{
		Name: "bootstrappah",
		Work: bootstrappah.Work,
	})

	s.Supervise(&supervisor.Worker{
		Name:     "invoker",
		Work:     invoker.Work,
		Requires: []string{"aggregator"},
	})

	s.Supervise(&supervisor.Worker{
		Name:     "api",
		Work:     apiServer.Work,
		Requires: []string{"invoker"},
	})

	if errs := s.Run(); len(errs) > 0 {
		for _, err := range errs {
			log.Println(err)
		}
		os.Exit(1)
	}
}

func main() {
	if err := rootCmd.Execute(); err != nil {
		log.Fatalln(err)
	}
}
<|endoftext|>


<commit_before>package libvirt

import (
	"encoding/xml"
	"fmt"
	"subuk/vmango/compute"
	"subuk/vmango/util"

	libvirt "github.com/libvirt/libvirt-go"
	libvirtxml "github.com/libvirt/libvirt-go-xml"
	"github.com/rs/zerolog"
)

type NodeRepository struct {
	pool   *ConnectionPool
	logger zerolog.Logger
}

func NewNodeRepository(pool *ConnectionPool, logger zerolog.Logger) *NodeRepository {
	return &NodeRepository{pool: pool, logger: logger}
}

func (repo *NodeRepository) List(options compute.NodeListOptions) ([]*compute.Node, error) {
	nodes := []*compute.Node{}
	for _, nodeId := range repo.pool.Nodes(nil) {
		node, err := repo.Get(nodeId, compute.NodeGetOptions{NoPins: options.NoPins})
		if err != nil {
			repo.logger.Warn().Err(err).Str("node", nodeId).Msg("cannot get node")
			continue
		}
		nodes = append(nodes, node)
	}
	return nodes, nil
}

func (repo *NodeRepository) Get(nodeId string, options compute.NodeGetOptions) (*compute.Node, error) {
	conn, err := repo.pool.Acquire(nodeId)
	if err != nil {
		return nil, util.NewError(err, "cannot acquire libvirt connection")
	}
	defer repo.pool.Release(nodeId)

	sysinfoXml, err := conn.GetSysinfo(0)
	if err != nil {
		return nil, util.NewError(err, "cannot get sysinfo")
	}
	sysinfoConfig := &libvirtxml.DomainSysInfo{}
	if err := xml.Unmarshal([]byte(sysinfoXml), sysinfoConfig); err != nil {
		return nil, util.NewError(err, "cannot parse sysinfo")
	}

	capsXml, err := conn.GetCapabilities()
	if err != nil {
		return nil, util.NewError(err, "cannot fetch host capabilities")
	}
	capsConfig := &libvirtxml.Caps{}
	if err := capsConfig.Unmarshal(capsXml); err != nil {
		return nil, util.NewError(err, "cannot parse capabilities")
	}
	hostname, err := conn.GetHostname()
	if err != nil {
		return nil, util.NewError(err, "cannot get hostname")
	}
	node := &compute.Node{
		Id:       nodeId,
		Hostname: hostname,
	}

	if len(sysinfoConfig.Processor) > 0 {
		for _, entry := range sysinfoConfig.Processor[0].Entry {
			if entry.Name == "version" {
				node.CpuInfo = entry.Value
			}
		}
	}
	if capsConfig.Host.IOMMU != nil {
		node.Iommu = capsConfig.Host.IOMMU.Support == "yes"
	}
	if capsConfig.Host.CPU != nil {
		node.CpuVendor = capsConfig.Host.CPU.Vendor
		node.CpuModel = capsConfig.Host.CPU.Model
		if capsConfig.Host.CPU.Topology != nil {
			node.ThreadsPerCore = capsConfig.Host.CPU.Topology.Threads
		}
	}
	switch capsConfig.Host.CPU.Arch {
	default:
		node.CpuArch = compute.ArchUnknown
	case "x86_64":
		node.CpuArch = compute.ArchAmd64
	}

	reqFreePagesMap := map[uint64]struct{}{}
	nodeCpuMap := map[int]compute.NodeCpu{}
	if capsConfig.Host.NUMA != nil && capsConfig.Host.NUMA.Cells != nil {
		node.Numas = make([]compute.NodeNuma, capsConfig.Host.NUMA.Cells.Num)
		for _, numaInfo := range capsConfig.Host.NUMA.Cells.Cells {
			if numaInfo.Memory != nil {
				node.Numas[numaInfo.ID].Memory = ComputeSizeFromLibvirtSize(numaInfo.Memory.Unit, numaInfo.Memory.Size)
			}
			for _, pageInfo := range numaInfo.PageInfo {
				switch pageInfo.Size {
				case 4:
					node.Numas[numaInfo.ID].Pages4k = pageInfo.Count
					reqFreePagesMap[4] = struct{}{}
				case 2048:
					node.Numas[numaInfo.ID].Pages2m = pageInfo.Count
					reqFreePagesMap[2048] = struct{}{}
				case 1048576:
					node.Numas[numaInfo.ID].Pages1g = pageInfo.Count
					reqFreePagesMap[1048576] = struct{}{}
				}
			}

			if numaInfo.CPUS != nil {
				for _, numaCpuInfo := range numaInfo.CPUS.CPUs {
					cpu := compute.NodeCpu{
						NumaId: numaInfo.ID,
					}
					if numaCpuInfo.SocketID != nil {
						cpu.SocketId = *numaCpuInfo.SocketID
					}
					if numaCpuInfo.CoreID != nil {
						cpu.CoreId = *numaCpuInfo.CoreID
					}
					nodeCpuMap[numaCpuInfo.ID] = cpu
				}
			}
		}
	}
	node.Cpus = make([]compute.NodeCpu, len(nodeCpuMap))
	for idx := 0; idx < len(node.Cpus); idx++ {
		node.Cpus[idx] = nodeCpuMap[idx]
	}

	reqFreePages := []uint64{}
	for pageSize := range reqFreePagesMap {
		reqFreePages = append(reqFreePages, pageSize)
	}
	freePages, err := conn.GetFreePages(reqFreePages, 0, uint(len(node.Numas)), 0)
	if err != nil {
		return nil, util.NewError(err, "cannot get free memory pages")
	}
	for idx, pageCount := range freePages {
		pageSize := reqFreePages[idx%len(reqFreePages)]
		numaId := idx / len(reqFreePages)
		switch pageSize {
		case 4:
			node.Numas[numaId].Pages4kFree = pageCount
		case 2048:
			node.Numas[numaId].Pages2mFree = pageCount
		case 1048576:
			node.Numas[numaId].Pages1gFree = pageCount
		}
	}

	if options.NoPins {
		return node, nil
	}

	virDomains, err := conn.ListAllDomains(0)
	if err != nil {
		return nil, util.NewError(err, "cannot list node domains")
	}
	for _, virDomain := range virDomains {
		virDomainXml, err := virDomain.GetXMLDesc(libvirt.DOMAIN_XML_INACTIVE)
		if err != nil {
			return nil, util.NewError(err, "cannot get domain xml")
		}
		virDomainConfig := &libvirtxml.Domain{}
		if err := virDomainConfig.Unmarshal(virDomainXml); err != nil {
			return nil, util.NewError(err, "cannot unmarshal domain xml")
		}
		if virDomainConfig.CPUTune == nil {
			for cpuId := range node.Cpus {
				node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
					Desc: "all",
					VmId: virDomainConfig.Name,
				})
			}
			continue
		}
		if virDomainConfig.CPUTune.VCPUPin == nil {
			for cpuId := range node.Cpus {
				node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
					Desc: "vcpus",
					VmId: virDomainConfig.Name,
				})
			}
			continue
		}

		if virDomainConfig.CPUTune.EmulatorPin == nil {
			for cpuId := range node.Cpus {
				node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
					Desc: "emulator",
					VmId: virDomainConfig.Name,
				})
			}
			continue
		}

		affinity := ParseCpuAffinity(virDomainConfig.CPUTune.EmulatorPin.CPUSet)
		for _, cpuId := range affinity {
			node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
				Desc: "emulator",
				VmId: virDomainConfig.Name,
			})
		}

		for _, vcpupin := range virDomainConfig.CPUTune.VCPUPin {
			vcpuAffinity := ParseCpuAffinity(vcpupin.CPUSet)
			for _, cpuId := range vcpuAffinity {
				node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
					Desc: fmt.Sprintf("vcpu-%d", vcpupin.VCPU),
					VmId: virDomainConfig.Name,
				})
			}
		}
	}
	if options.CpuNumaIdFilter {
		newCpus := []compute.NodeCpu{}
		for _, cpu := range node.Cpus {
			if cpu.NumaId == options.CpuNumaId {
				newCpus = append(newCpus, cpu)
			} else {
				newCpus = append(newCpus, compute.NodeCpu{SocketId: -1, CoreId: -1, NumaId: -1})
			}
		}
		node.Cpus = newCpus
	}
	return node, nil
}
<commit_msg>Query memory pages info only if hypervisor supports it<commit_after>package libvirt

import (
	"encoding/xml"
	"fmt"
	"subuk/vmango/compute"
	"subuk/vmango/util"

	libvirt "github.com/libvirt/libvirt-go"
	libvirtxml "github.com/libvirt/libvirt-go-xml"
	"github.com/rs/zerolog"
)

type NodeRepository struct {
	pool   *ConnectionPool
	logger zerolog.Logger
}

func NewNodeRepository(pool *ConnectionPool, logger zerolog.Logger) *NodeRepository {
	return &NodeRepository{pool: pool, logger: logger}
}

func (repo *NodeRepository) List(options compute.NodeListOptions) ([]*compute.Node, error) {
	nodes := []*compute.Node{}
	for _, nodeId := range repo.pool.Nodes(nil) {
		node, err := repo.Get(nodeId, compute.NodeGetOptions{NoPins: options.NoPins})
		if err != nil {
			repo.logger.Warn().Err(err).Str("node", nodeId).Msg("cannot get node")
			continue
		}
		nodes = append(nodes, node)
	}
	return nodes, nil
}

func (repo *NodeRepository) Get(nodeId string, options compute.NodeGetOptions) (*compute.Node, error) {
	conn, err := repo.pool.Acquire(nodeId)
	if err != nil {
		return nil, util.NewError(err, "cannot acquire libvirt connection")
	}
	defer repo.pool.Release(nodeId)

	sysinfoXml, err := conn.GetSysinfo(0)
	if err != nil {
		return nil, util.NewError(err, "cannot get sysinfo")
	}
	sysinfoConfig := &libvirtxml.DomainSysInfo{}
	if err := xml.Unmarshal([]byte(sysinfoXml), sysinfoConfig); err != nil {
		return nil, util.NewError(err, "cannot parse sysinfo")
	}

	capsXml, err := conn.GetCapabilities()
	if err != nil {
		return nil, util.NewError(err, "cannot fetch host capabilities")
	}
	capsConfig := &libvirtxml.Caps{}
	if err := capsConfig.Unmarshal(capsXml); err != nil {
		return nil, util.NewError(err, "cannot parse capabilities")
	}
	hostname, err := conn.GetHostname()
	if err != nil {
		return nil, util.NewError(err, "cannot get hostname")
	}
	node := &compute.Node{
		Id:       nodeId,
		Hostname: hostname,
	}

	if len(sysinfoConfig.Processor) > 0 {
		for _, entry := range sysinfoConfig.Processor[0].Entry {
			if entry.Name == "version" {
				node.CpuInfo = entry.Value
			}
		}
	}
	if capsConfig.Host.IOMMU != nil {
		node.Iommu = capsConfig.Host.IOMMU.Support == "yes"
	}
	if capsConfig.Host.CPU != nil {
		node.CpuVendor = capsConfig.Host.CPU.Vendor
		node.CpuModel = capsConfig.Host.CPU.Model
		if capsConfig.Host.CPU.Topology != nil {
			node.ThreadsPerCore = capsConfig.Host.CPU.Topology.Threads
		}
	}
	switch capsConfig.Host.CPU.Arch {
	default:
		node.CpuArch = compute.ArchUnknown
	case "x86_64":
		node.CpuArch = compute.ArchAmd64
	}

	reqFreePagesMap := map[uint64]struct{}{}
	nodeCpuMap := map[int]compute.NodeCpu{}
	if capsConfig.Host.NUMA != nil && capsConfig.Host.NUMA.Cells != nil {
		node.Numas = make([]compute.NodeNuma, capsConfig.Host.NUMA.Cells.Num)
		for _, numaInfo := range capsConfig.Host.NUMA.Cells.Cells {
			if numaInfo.Memory != nil {
				node.Numas[numaInfo.ID].Memory = ComputeSizeFromLibvirtSize(numaInfo.Memory.Unit, numaInfo.Memory.Size)
			}
			for _, pageInfo := range numaInfo.PageInfo {
				switch pageInfo.Size {
				case 4:
					node.Numas[numaInfo.ID].Pages4k = pageInfo.Count
					reqFreePagesMap[4] = struct{}{}
				case 2048:
					node.Numas[numaInfo.ID].Pages2m = pageInfo.Count
					reqFreePagesMap[2048] = struct{}{}
				case 1048576:
					node.Numas[numaInfo.ID].Pages1g = pageInfo.Count
					reqFreePagesMap[1048576] = struct{}{}
				}
			}

			if numaInfo.CPUS != nil {
				for _, numaCpuInfo := range numaInfo.CPUS.CPUs {
					cpu := compute.NodeCpu{
						NumaId: numaInfo.ID,
					}
					if numaCpuInfo.SocketID != nil {
						cpu.SocketId = *numaCpuInfo.SocketID
					}
					if numaCpuInfo.CoreID != nil {
						cpu.CoreId = *numaCpuInfo.CoreID
					}
					nodeCpuMap[numaCpuInfo.ID] = cpu
				}
			}
		}
	}
	node.Cpus = make([]compute.NodeCpu, len(nodeCpuMap))
	for idx := 0; idx < len(node.Cpus); idx++ {
		node.Cpus[idx] = nodeCpuMap[idx]
	}

	if len(reqFreePagesMap) > 0 {
		reqFreePages := []uint64{}
		for pageSize := range reqFreePagesMap {
			reqFreePages = append(reqFreePages, pageSize)
		}
		freePages, err := conn.GetFreePages(reqFreePages, 0, uint(len(node.Numas)), 0)
		if err != nil {
			return nil, util.NewError(err, "cannot get free memory pages")
		}
		for idx, pageCount := range freePages {
			pageSize := reqFreePages[idx%len(reqFreePages)]
			numaId := idx / len(reqFreePages)
			switch pageSize {
			case 4:
				node.Numas[numaId].Pages4kFree = pageCount
			case 2048:
				node.Numas[numaId].Pages2mFree = pageCount
			case 1048576:
				node.Numas[numaId].Pages1gFree = pageCount
			}
		}
	}

	if options.NoPins {
		return node, nil
	}

	virDomains, err := conn.ListAllDomains(0)
	if err != nil {
		return nil, util.NewError(err, "cannot list node domains")
	}
	for _, virDomain := range virDomains {
		virDomainXml, err := virDomain.GetXMLDesc(libvirt.DOMAIN_XML_INACTIVE)
		if err != nil {
			return nil, util.NewError(err, "cannot get domain xml")
		}
		virDomainConfig := &libvirtxml.Domain{}
		if err := virDomainConfig.Unmarshal(virDomainXml); err != nil {
			return nil, util.NewError(err, "cannot unmarshal domain xml")
		}
		if virDomainConfig.CPUTune == nil {
			for cpuId := range node.Cpus {
				node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
					Desc: "all",
					VmId: virDomainConfig.Name,
				})
			}
			continue
		}
		if virDomainConfig.CPUTune.VCPUPin == nil {
			for cpuId := range node.Cpus {
				node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
					Desc: "vcpus",
					VmId: virDomainConfig.Name,
				})
			}
			continue
		}

		if virDomainConfig.CPUTune.EmulatorPin == nil {
			for cpuId := range node.Cpus {
				node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
					Desc: "emulator",
					VmId: virDomainConfig.Name,
				})
			}
			continue
		}

		affinity := ParseCpuAffinity(virDomainConfig.CPUTune.EmulatorPin.CPUSet)
		for _, cpuId := range affinity {
			node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
				Desc: "emulator",
				VmId: virDomainConfig.Name,
			})
		}

		for _, vcpupin := range virDomainConfig.CPUTune.VCPUPin {
			vcpuAffinity := ParseCpuAffinity(vcpupin.CPUSet)
			for _, cpuId := range vcpuAffinity {
				node.Cpus[cpuId].Pins = append(node.Cpus[cpuId].Pins, compute.NodeCpuPin{
					Desc: fmt.Sprintf("vcpu-%d", vcpupin.VCPU),
					VmId: virDomainConfig.Name,
				})
			}
		}
	}
	if options.CpuNumaIdFilter {
		newCpus := []compute.NodeCpu{}
		for _, cpu := range node.Cpus {
			if cpu.NumaId == options.CpuNumaId {
				newCpus = append(newCpus, cpu)
			} else {
				newCpus = append(newCpus, compute.NodeCpu{SocketId: -1, CoreId: -1, NumaId: -1})
			}
		}
		node.Cpus = newCpus
	}
	return node, nil
}
<|endoftext|>


<commit_before>package midi

import (
	"sync"
	"time"

	"github.com/emicklei/melrose/core"
	"github.com/emicklei/melrose/notify"
)

type KeyTrigger struct {
	mutex   *sync.RWMutex
	playing bool
	ctx     core.Context
	channel int
	note    core.Note
	fun     core.Valueable
}

func NewKeyTrigger(ctx core.Context, channel int, onNote core.Note, startStop core.Valueable) *KeyTrigger {
	return &KeyTrigger{
		mutex:   new(sync.RWMutex),
		ctx:     ctx,
		channel: channel,
		note:    onNote,
		fun:     startStop}
}

// NoteOn is part of core.NoteListener
func (t *KeyTrigger) NoteOn(channel int, n core.Note) {
	if core.IsDebug() {
		notify.Debugf("keytrigger.NoteOn ch=%d note=%v", channel, n)
	}
	if channel != t.channel {
		return
	}
	if n.Name != t.note.Name {
		return
	}
	if n.Octave != t.note.Octave {
		return
	}
	val := t.fun.Value()
	if val == nil {
		return
	}
	t.mutex.Lock()
	defer t.mutex.Unlock()
	// both playable and evaluatable are allowed
	if play, ok := val.(core.Playable); ok {
		stopper, stoppable := val.(core.Stoppable)
		if stoppable {
			if t.playing {
				stopper.Stop(t.ctx)
				t.playing = false
			} else {
				t.playing = true
				_ = play.Play(t.ctx, time.Now())
			}
			return
		}
		// cannot stop
		_ = play.Play(t.ctx, time.Now())
		return
	}
	// not playable, maybe evaluatable
	if eval, ok := val.(core.Evaluatable); ok {
		eval.Evaluate(t.ctx)
	}
}

// NoteOff is part of core.NoteListener
func (t *KeyTrigger) NoteOff(channel int, n core.Note) {
	if core.IsDebug() {
		notify.Debugf("keytrigger.NoteOff ch=%d note=%v", channel, n)
	}
	// key trigger is not interested in this
}

func (t *KeyTrigger) ControlChange(channel, number, value int) {
	if core.IsDebug() {
		notify.Debugf("keytrigger.ControlChange %d %d %d", channel, number, value)
	}
	// key trigger is not interested in this
}
<commit_msg>show key hit feedback<commit_after>package midi

import (
	"sync"
	"time"

	"github.com/emicklei/melrose/core"
	"github.com/emicklei/melrose/notify"
)

type KeyTrigger struct {
	mutex   *sync.RWMutex
	playing bool
	ctx     core.Context
	channel int
	note    core.Note
	fun     core.Valueable
}

func NewKeyTrigger(ctx core.Context, channel int, onNote core.Note, startStop core.Valueable) *KeyTrigger {
	return &KeyTrigger{
		mutex:   new(sync.RWMutex),
		ctx:     ctx,
		channel: channel,
		note:    onNote,
		fun:     startStop}
}

// NoteOn is part of core.NoteListener
func (t *KeyTrigger) NoteOn(channel int, n core.Note) {
	if core.IsDebug() {
		notify.Debugf("keytrigger.NoteOn ch=%d note=%v", channel, n)
	}
	if channel != t.channel {
		return
	}
	if n.Name != t.note.Name {
		return
	}
	if n.Octave != t.note.Octave {
		return
	}
	val := t.fun.Value()
	if val == nil {
		return
	}
	t.mutex.Lock()
	defer t.mutex.Unlock()
	// both playable and evaluatable are allowed
	if play, ok := val.(core.Playable); ok {
		stopper, stoppable := val.(core.Stoppable)
		if stoppable {
			if t.playing {
				notify.Infof("%s -> stop(%s)", t.note.String(), core.Storex(t.fun))
				stopper.Stop(t.ctx)
				t.playing = false
			} else {
				t.playing = true
				notify.Infof("%s -> play(%s)", t.note.String(), core.Storex(t.fun))
				_ = play.Play(t.ctx, time.Now())
			}
			return
		}
		// cannot stop
		_ = play.Play(t.ctx, time.Now())
		return
	}
	// not playable, maybe evaluatable
	if eval, ok := val.(core.Evaluatable); ok {
		eval.Evaluate(t.ctx)
	}
}

// NoteOff is part of core.NoteListener
func (t *KeyTrigger) NoteOff(channel int, n core.Note) {
	if core.IsDebug() {
		notify.Debugf("keytrigger.NoteOff ch=%d note=%v", channel, n)
	}
	// key trigger is not interested in this
}

func (t *KeyTrigger) ControlChange(channel, number, value int) {
	if core.IsDebug() {
		notify.Debugf("keytrigger.ControlChange %d %d %d", channel, number, value)
	}
	// key trigger is not interested in this
}
<|endoftext|>


<commit_before>package sudoku

// Holds the global Sudoku board state

// Communication is handled by four major mechanics:
// * A channel to distribute the current board state
// * A channel to notify threads of updates
// * A channel to update known & possible values

import "math"

const (
	boardRow    = 0
	boardCol    = 1
	boardSquare = 2
)

// coord contains x and y elements for a given position
type coord struct {
	x int
	y int
}

// A cell holds all of the required knowledge for a cell.
// It contains it's own address on the board (guarenteed unique)
// if actual is set to 0, actual value is unknown
type cell struct {
	location coord
	actual   int
	possible []int
}

type cluster []cell

type board struct {
	size     int
	clusters []cluster
}

/*
type Cluster interface {
	Len() int
	Swap(i, j int)
	Less(i, j int)
}
*/

func (c Cluster) Len() int {
	return len(c)
}
func (c Cluster) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}
func (c Cluster) Less(i, j int) bool {
	if c[i].x < c[j].x {
		return true
	}
	if c[i].x == c[j].x && c[i].y < c[j].y {
		return true
	}
	return false
}

func createBoard(size int) board {
	var newBoard board
	newBoard.size = size
	for i := 0; i < size*size; i++ {
		for j := 0; j < size*size; j++ {
			board[coord{x: i, y: j}].location = coord{x: i, y: j}
			for k := 1; k <= size*size; k++ {
				board[coord{x: i, y: j}].possible = append(board[coord{x: i, y: j}].possible, k)
			}
		}
	}
	return newBoard
}

func getPos(position coord, orientation, size int) (int, error) {
	if position.x >= size*size {
		return -1, errors.New("x position is larger than the board")
	}
	if position.y >= size*size {
		return -1, errors.New("y position is larger than the board")
	}
	switch orientation {
	case boardRow:
		return position.x
	case boardCol:
		return position.y
	case boardSquare:
		return ((position.y / size) * size) + (position.x / size)
	default:
		return -1, errors.New("bad position")
	}
}

func clusterPicker(in board, orient int, position coord) (cluster, error) {
	if position.x >= len(in) {
		return cluster{}, errors.New("x coord out of range")
	} else if position.y >= len(in) {
		return cluster{}, errors.New("y coord out of range")
	}
	switch orient {
	case boardRow:
		return in[position.y], nil
	case boardCol:
		var result cluster
		for _, each := range in {
			result = append(result, each[position.x])
		}
		return outCluster, nil
	case boardSquare:
		var result cluster
		boardSize := int(math.Sqrt(len(board)))
		startX := (position.x / boardSize) * boardSize
		endX := ((position.x / boardSize) + 1) * boardSize
		startY := (position.y / boardSize) * boardSize
		endY := ((position.y / boardSize) + 1) * boardSize
		for x := startX; x < endX; x++ {
			for y := startY; y < endY; y++ {
				result = append(result, in[x][y])
			}
		}
		return result, nil
	default:
		return cluster{}, errors.New("bad orientation")
	}
}

func clusterFilter(update <-chan coord, in <-chan board, out [][]chan<- cluster, status [][]<-chan struct{}) {

	var toWork cluster
	defer closeArrArrChan(out)
	// don't actually do anything until you have a board state to work with
	<-in
	for {
		select {
		case changed, more := <-update:
			if !more {
				return
			}
			for i = boardRow; i <= boardSquare; i++ {
				curBoard := <-in
				position := getPos(changed, i, curBoard.size)
				if _, open := <-status[i][position]; !open {
					// skip this update if the cluster is already solved
					continue
				}
				curCluster, err := clusterPicker(curBoard, i, changed)
				if err != nil {
					panic(err) // #TODO# replace this panic
				}
				out[i][position] <- curCluster
			}
		}
	}
}

// takes an input and sends it out as it can
// exits when in is closed
// closes out on exit
func clusterSticky(in <-chan cluster, out chan<- cluster) {
	var more bool
	defer close(out)

	// preload the locCluster - don't do anything until you have that
	locCluster := <-in
	for {
		select {
		case locCluster, more = <-in:
			// if you get an update from upstream, do that first
			if !more {
				return
			}
		// pass the update downstream then block till you get another update from upstream
		case out <- locCluster:
			locCluster = <-in
		}
	}
}

// like a buffered channel, but no limit to the buffer size
// closes out on exit
// exits when in is closed
func updateBuffer(in <-chan cell, out chan<- cell) {
	var updates []interface{}
	var singleUpdate interface{}
	var open bool
	defer close(out)

	for {
		if len(updates) < 1 {
			// if you currently don't have anything to pass, WAIT FOR SOMETHING
			singleUpdate, open = <-in
			if !open {
				return
			}
			updates = append(updates, singleUpdate)
			continue
		}
		select {
		case singleUpdate, open = <-in:
			if !open {
				return
			}
			updates = append(updates, singleUpdate)
		case out <- updates[0]:
			updates = updates[1:]
		}
	}
}

// boardCache serves a given (or newer) update out as many times as requested
// closes `out` on exit
// exits when in `is` closed
func boardCache(in chan board, out chan<- board) {
	currentBoard := <-in
	var done bool
	defer close(out)

	for {
		select {
		case currentBoard, done = <-in:
			if done {
				return
			}
		case in <- currentBoard:
		case out <- currentBoard:
		}
	}
}

// looks for a send from all status channels
// if recieved or closed on all status chan, send on `idle`
// if all status is closed, close idle and exits
func idleCheck(status [][]<-chan interface{}, idle chan<- interface{}) {
	var more, isIdle, isSolved bool
	defer close(idle)

	for {
	start:
		isSolved, isIdle = true, true
		for _, middle := range status {
			for _, inner := range status {
				select {
				case _, more = <-inner:
					if more {
						isSolved = false
					}
				default:
					isSolved = false
					isIdle = false
					continue start
				}
			}
		}
		if isSolved {
			return
		} else if isIdle {
			idle <- nil
		}
	}
}

// takes a given cluster, and runs it through all of the moves
// exits when one of the conditions is met, or when the update channel is closed
// closes the status channel on exit
func clusterWorker(in <-chan cluster, status chan<- struct{}, updates chan<- cell, problems chan<- error) {
	defer close(status)

	var more bool
	var newCluster cluster
	var index indexedCluster
	var changes, newChanges []cell

	for {
		select {
		case newCluster, more = <-in:
			if !more {
				// if the channel is closed, exit
				return
			}
			if clusterSolved(cluster) {
				// if the cell is solved, exit
				return
			}

			newChanges = solvedNoPossible(cluster.cells)
			changes = append(changes, newChanges)

			newChanges = eliminateKnowns(cluster.cells)
			changes = append(changes, newChanges)

			newChanges = singleValueSolver(cluster.cells)
			changes = append(changes, newChanges)

			newChanges = cellLimiter(cluster.cells)
			changes = append(changes, newChanges)

			index = indexCluster(cluster.cells)

			newChanges = singleCellSolver(index, cluster.cells)
			changes = append(changes, newChanges)

			newChanges = valueLimiter(index, cluster.cells)
			changes = append(changes, newChanges)

			// feed all those changes into the update queue
			for len(changes) > 1 {
				updates <- changes[0]
				changes = changes[1:]
			}
		case status <- nil:
			// report idle only if there is nothing to do - order matters
		}
	}
}

//
func updateProcessor(curBoard chan board, status <-chan struct{}, updates <-chan cell, posChange chan<- coord, problems <-chan error) {
	defer close(curBoard)
	defer close(posChange)

}
<commit_msg>started writing a function to change the board as needed<commit_after>package sudoku

// Holds the global Sudoku board state

// Communication is handled by four major mechanics:
// * A channel to distribute the current board state
// * A channel to notify threads of updates
// * A channel to update known & possible values

import "math"

const (
	boardRow    = 0
	boardCol    = 1
	boardSquare = 2
)

// coord contains x and y elements for a given position
type coord struct {
	x int
	y int
}

// A cell holds all of the required knowledge for a cell.
// It contains it's own address on the board (guarenteed unique)
// if actual is set to 0, actual value is unknown
type cell struct {
	location coord
	actual   int
	possible []int
	excluded []int
}

type cluster []cell

type board struct {
	size     int
	clusters []cluster
}

/*
type Cluster interface {
	Len() int
	Swap(i, j int)
	Less(i, j int)
}
*/

func (c Cluster) Len() int {
	return len(c)
}
func (c Cluster) Swap(i, j int) {
	c[i], c[j] = c[j], c[i]
}
func (c Cluster) Less(i, j int) bool {
	if c[i].x < c[j].x {
		return true
	}
	if c[i].x == c[j].x && c[i].y < c[j].y {
		return true
	}
	return false
}

func createBoard(size int) board {
	var newBoard board
	newBoard.size = size
	for i := 0; i < size*size; i++ {
		for j := 0; j < size*size; j++ {
			board[coord{x: i, y: j}].location = coord{x: i, y: j}
			for k := 1; k <= size*size; k++ {
				board[coord{x: i, y: j}].possible = append(board[coord{x: i, y: j}].possible, k)
			}
		}
	}
	return newBoard
}

func getPos(position coord, orientation, size int) (int, error) {
	if position.x >= size*size {
		return -1, errors.New("x position is larger than the board")
	}
	if position.y >= size*size {
		return -1, errors.New("y position is larger than the board")
	}
	switch orientation {
	case boardRow:
		return position.x
	case boardCol:
		return position.y
	case boardSquare:
		return ((position.y / size) * size) + (position.x / size)
	default:
		return -1, errors.New("bad position")
	}
}

func clusterPicker(in board, orient int, position coord) (cluster, error) {
	if position.x >= len(in) {
		return cluster{}, errors.New("x coord out of range")
	} else if position.y >= len(in) {
		return cluster{}, errors.New("y coord out of range")
	}
	switch orient {
	case boardRow:
		return in[position.y], nil
	case boardCol:
		var result cluster
		for _, each := range in {
			result = append(result, each[position.x])
		}
		return outCluster, nil
	case boardSquare:
		var result cluster
		boardSize := int(math.Sqrt(len(board)))
		startX := (position.x / boardSize) * boardSize
		endX := ((position.x / boardSize) + 1) * boardSize
		startY := (position.y / boardSize) * boardSize
		endY := ((position.y / boardSize) + 1) * boardSize
		for x := startX; x < endX; x++ {
			for y := startY; y < endY; y++ {
				result = append(result, in[x][y])
			}
		}
		return result, nil
	default:
		return cluster{}, errors.New("bad orientation")
	}
}

func clusterFilter(update <-chan coord, in <-chan board, out [][]chan<- cluster, status [][]<-chan struct{}) {

	var toWork cluster
	defer closeArrArrChan(out)
	// don't actually do anything until you have a board state to work with
	<-in
	for {
		select {
		case changed, more := <-update:
			if !more {
				return
			}
			for i = boardRow; i <= boardSquare; i++ {
				curBoard := <-in
				position := getPos(changed, i, curBoard.size)
				if _, open := <-status[i][position]; !open {
					// skip this update if the cluster is already solved
					continue
				}
				curCluster, err := clusterPicker(curBoard, i, changed)
				if err != nil {
					panic(err) // #TODO# replace this panic
				}
				out[i][position] <- curCluster
			}
		}
	}
}

// takes an input and sends it out as it can
// exits when in is closed
// closes out on exit
func clusterSticky(in <-chan cluster, out chan<- cluster) {
	var more bool
	defer close(out)

	// preload the locCluster - don't do anything until you have that
	locCluster := <-in
	for {
		select {
		case locCluster, more = <-in:
			// if you get an update from upstream, do that first
			if !more {
				return
			}
		// pass the update downstream then block till you get another update from upstream
		case out <- locCluster:
			locCluster = <-in
		}
	}
}

// like a buffered channel, but no limit to the buffer size
// closes out on exit
// exits when in is closed
func updateBuffer(in <-chan cell, out chan<- cell) {
	var updates []interface{}
	var singleUpdate interface{}
	var open bool
	defer close(out)

	for {
		if len(updates) < 1 {
			// if you currently don't have anything to pass, WAIT FOR SOMETHING
			singleUpdate, open = <-in
			if !open {
				return
			}
			updates = append(updates, singleUpdate)
			continue
		}
		select {
		case singleUpdate, open = <-in:
			if !open {
				return
			}
			updates = append(updates, singleUpdate)
		case out <- updates[0]:
			updates = updates[1:]
		}
	}
}

// boardCache serves a given (or newer) update out as many times as requested
// closes `out` on exit
// exits when in `is` closed
func boardCache(in chan board, out chan<- board) {
	currentBoard := <-in
	var done bool
	defer close(out)

	for {
		select {
		case currentBoard, done = <-in:
			if done {
				return
			}
		case in <- currentBoard:
		case out <- currentBoard:
		}
	}
}

// looks for a send from all status channels
// if recieved or closed on all status chan, send on `idle`
// if all status is closed, close idle and exits
func idleCheck(status [][]<-chan interface{}, idle chan<- interface{}) {
	var more, isIdle, isSolved bool
	defer close(idle)

	for {
	start:
		isSolved, isIdle = true, true
		for _, middle := range status {
			for _, inner := range status {
				select {
				case _, more = <-inner:
					if more {
						isSolved = false
					}
				default:
					isSolved = false
					isIdle = false
					continue start
				}
			}
		}
		if isSolved {
			return
		} else if isIdle {
			idle <- nil
		}
	}
}

// takes a given cluster, and runs it through all of the moves
// exits when one of the conditions is met, or when the update channel is closed
// closes the status channel on exit
func clusterWorker(in <-chan cluster, status chan<- struct{}, updates chan<- cell, problems chan<- error) {
	defer close(status)

	var more bool
	var newCluster cluster
	var index indexedCluster
	var changes, newChanges []cell

	for {
		select {
		case newCluster, more = <-in:
			if !more {
				// if the channel is closed, exit
				return
			}
			if clusterSolved(cluster) {
				// if the cell is solved, exit
				return
			}

			newChanges = solvedNoPossible(cluster.cells)
			changes = append(changes, newChanges)

			newChanges = eliminateKnowns(cluster.cells)
			changes = append(changes, newChanges)

			newChanges = singleValueSolver(cluster.cells)
			changes = append(changes, newChanges)

			newChanges = cellLimiter(cluster.cells)
			changes = append(changes, newChanges)

			index = indexCluster(cluster.cells)

			newChanges = singleCellSolver(index, cluster.cells)
			changes = append(changes, newChanges)

			newChanges = valueLimiter(index, cluster.cells)
			changes = append(changes, newChanges)

			// feed all those changes into the update queue
			for len(changes) > 1 {
				updates <- changes[0]
				changes = changes[1:]
			}
		case status <- nil:
			// report idle only if there is nothing to do - order matters
		}
	}
}

// processes updates
// priority is problems, updates, then status checks
func updateProcessor(curBoard chan board, status <-chan struct{}, updates <-chan cell, posChange chan<- coord, problems <-chan error) {
	defer close(curBoard)
	defer close(posChange)

	var newBoard board
	var err error

	for {
		select {
		case err = <-problems:
			// any errors get top priority
		case cellChange = <-updates:
			// any udpates are handled before idle checks
			newBoard, err = changeBoard(<-curBoard, cellChange)
		case _, solved = <-status:
			// status check if neither of the above happens - or solved check
		}
	}
}

func changeBoard(in board, u cell) (out board, err error) {
	target := in.clusters[u.location.x][u.location.y]
	if cell.actual != 0 {
		if target.actual != 0
		for i = 1; i <= board.level; i++ {
			cell.excluded = append(cell.excluded, i)
		}
		cell.excluded = dedupArr(cell.excluded)
		in.clusters[cell.location.x][cell.location.y] = cell
	} else {
		in.clusters[cell.location.x][cell.location.y].excluded = addArr(board.clusters[cell.location.x][cell.location.y].excluded, cell.excluded)
	}

}
<|endoftext|>


<commit_before>/*
Copyright 2017 Tuenti Technologies S.L. All rights reserved.

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 pouch

import (
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"time"

	"github.com/hashicorp/vault/api"
)

const (
	DefaultStatePath           = "/var/lib/pouch/state"
	DefaultStateMode           = os.FileMode(0600)
	DefaultStateDirMode        = os.FileMode(0700)
	DefaultSecretDurationRatio = 0.75

	PreviousStateFilePostfix = "-prev"
)

type PouchState struct {
	// Last known token
	Token string `json:"token,omitempty"`

	// Secrets state
	Secrets map[string]*SecretState `json:"secrets,omitempty"`

	// Path from where this state was read
	Path string `json:"-"`
}

func NewState(path string) *PouchState {
	return &PouchState{Path: path}
}

func LoadState(path string) (*PouchState, error) {
	if path == "" {
		path = DefaultStatePath
	}
	d, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var state PouchState
	err = json.Unmarshal(d, &state)
	if err != nil {
		return nil, err
	}
	state.Path = path
	return &state, nil
}

func (s *PouchState) Save() error {
	path := s.Path
	if path == "" {
		path = DefaultStatePath
	}

	// Create directories if they don't exist
	dir := filepath.Dir(path)
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		err = os.MkdirAll(dir, DefaultStateDirMode)
		if err != nil {
			return err
		}
	}

	// Copy existing data to reduce risk of corrupting state
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		d, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}
		err = ioutil.WriteFile(path+PreviousStateFilePostfix, d, DefaultStateMode)
		if err != nil {
			return err
		}
	}

	// Finally write the state
	d, err := json.MarshalIndent(s, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(path, d, DefaultStateMode)
}

// Sources of TTUs
var secretTTUSources = []func(*SecretState) (*time.Time, error){
	ttuFromTTLOrLeaseDuration,
	ttuFromCertificateValidity,
}

func ttuFromTTLOrLeaseDuration(s *SecretState) (*time.Time, error) {
	ttl, ttlKnown := s.TTL()

	var duration int
	switch {
	case ttlKnown && s.LeaseDuration > 0:
		if ttl < s.LeaseDuration {
			duration = ttl
		} else {
			duration = s.LeaseDuration
		}
	case ttlKnown:
		duration = ttl
	case s.LeaseDuration > 0:
		duration = s.LeaseDuration
	default:
		return nil, nil
	}

	ttu := s.Timestamp.Add(time.Duration(float64(duration)*s.Ratio()) * time.Second)
	return &ttu, nil
}

func ttuFromCertificateValidity(s *SecretState) (*time.Time, error) {
	if s.Data == nil {
		return nil, nil
	}

	data, ok := s.Data["certificate"].(string)
	if !ok {
		return nil, nil
	}

	block, _ := pem.Decode([]byte(data))
	if block == nil {
		return nil, fmt.Errorf("failed to parse certificate PEM")
	}

	certificate, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, err
	}

	ttl := certificate.NotAfter.Sub(certificate.NotBefore)
	ttu := certificate.NotBefore.Add(time.Duration(float64(ttl) * s.Ratio()))
	return &ttu, nil
}

func (s *PouchState) SetSecret(name string, secret *api.Secret) {
	if s.Secrets == nil {
		s.Secrets = make(map[string]*SecretState)
	}
	state := &SecretState{
		Name:          name,
		Timestamp:     time.Now(),
		LeaseDuration: secret.LeaseDuration,
		Data:          secret.Data,
	}

	if _, known := state.TimeToUpdate(); !known {
		// Without a known TTU, we don't know when to update
		state.DisableAutoUpdate = true
	}

	if oldState, found := s.Secrets[name]; found {
		state.FilesUsing = oldState.FilesUsing
	}
	s.Secrets[name] = state
}

func (s *PouchState) DeleteSecret(name string) {
	delete(s.Secrets, name)
}

func (s *PouchState) NextUpdate() (secret *SecretState, minTTU time.Time) {
	for name := range s.Secrets {
		if s.Secrets[name].DisableAutoUpdate {
			continue
		}
		ttu, known := s.Secrets[name].TimeToUpdate()
		if known && (secret == nil || ttu.Before(minTTU)) {
			secret = s.Secrets[name]
			minTTU = ttu
		}
	}
	return
}

type PriorityFileSortedList []PriorityFile

type PriorityFile struct {
	Priority int    `json:"-"`
	Path     string `json:"path,omitempty"`
}

func (pf *PriorityFile) MarshalJSON() ([]byte, error) {
	return json.Marshal(pf.Path)
}

func (s *PriorityFileSortedList) UnmarshalJSON(data []byte) error {
	var priorityFiles []string

	if err := json.Unmarshal(data, &priorityFiles); err != nil {
		return err
	}

	// To keep the same state as when it was written, each file is assigned a
	// priority according to the order that they appear in the state file,
	for i, pf := range priorityFiles {
		*s = append(*s, PriorityFile{Path: pf, Priority: i * 10})
	}

	return nil
}

func (p PriorityFileSortedList) Len() int      { return len(p) }
func (p PriorityFileSortedList) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p PriorityFileSortedList) Less(i, j int) bool {
	if p[i].Priority != p[j].Priority {
		return p[i].Priority < p[j].Priority
	}
	return p[i].Path < p[j].Path
}

type SecretData map[string]interface{}

type SecretState struct {
	// Secret name
	Name string `json:"name,omitempty"`

	// Time when the secret was read
	Timestamp time.Time `json:"creation_time,omitempty"`

	// Lease duration, in seconds, if any when the secret was read
	LeaseDuration int `json:"lease_duration,omitempty"`

	// Secret will be renewed after this portion of its life has passed
	DurationRatio float64 `json:"duration_ratio,omitempty"`

	// If the secret has no expiration data, don't try to update it
	DisableAutoUpdate bool `json:"disable_auto_uptdate,omitempty"`

	// Actual secret
	Data SecretData `json:"data,omitempty"`

	// Files using this secret
	FilesUsing PriorityFileSortedList `json:"files_using,omitempty"`
}

func (s *SecretState) Ratio() float64 {
	ratio := s.DurationRatio
	if ratio == 0 {
		ratio = DefaultSecretDurationRatio
	}
	return ratio
}

func (s *SecretState) TTL() (int, bool) {
	if s.Data == nil {
		return 0, false
	}
	switch ttlNumber := s.Data["ttl"].(type) {
	case json.Number:
		ttl, err := ttlNumber.Int64()
		if err != nil {
			return 0, false
		}
		return int(ttl), true
	case int:
		return ttlNumber, true
	case int64:
		return int(ttlNumber), true
	}
	return 0, false
}

func (s *SecretState) TimeToUpdate() (time.Time, bool) {
	for _, source := range secretTTUSources {
		ttu, err := source(s)
		if err != nil {
			log.Printf("Error trying to obtain TTU for secret '%s': %s", s.Name, err)
			continue
		}
		if ttu != nil {
			return *ttu, true
		}
	}
	return time.Time{}, false
}

func (s *SecretState) RegisterUsage(path string, priority int) {
	for _, f := range s.FilesUsing {
		if f.Path == path {
			// Already registered
			return
		}
	}
	s.FilesUsing = append(s.FilesUsing, PriorityFile{Priority: priority, Path: path})
	sort.Sort(s.FilesUsing)
}
<commit_msg>Update each secret on minimum TTU found<commit_after>/*
Copyright 2017 Tuenti Technologies S.L. All rights reserved.

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 pouch

import (
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"sort"
	"time"

	"github.com/hashicorp/vault/api"
)

const (
	DefaultStatePath           = "/var/lib/pouch/state"
	DefaultStateMode           = os.FileMode(0600)
	DefaultStateDirMode        = os.FileMode(0700)
	DefaultSecretDurationRatio = 0.75

	PreviousStateFilePostfix = "-prev"
)

type PouchState struct {
	// Last known token
	Token string `json:"token,omitempty"`

	// Secrets state
	Secrets map[string]*SecretState `json:"secrets,omitempty"`

	// Path from where this state was read
	Path string `json:"-"`
}

func NewState(path string) *PouchState {
	return &PouchState{Path: path}
}

func LoadState(path string) (*PouchState, error) {
	if path == "" {
		path = DefaultStatePath
	}
	d, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var state PouchState
	err = json.Unmarshal(d, &state)
	if err != nil {
		return nil, err
	}
	state.Path = path
	return &state, nil
}

func (s *PouchState) Save() error {
	path := s.Path
	if path == "" {
		path = DefaultStatePath
	}

	// Create directories if they don't exist
	dir := filepath.Dir(path)
	if _, err := os.Stat(dir); os.IsNotExist(err) {
		err = os.MkdirAll(dir, DefaultStateDirMode)
		if err != nil {
			return err
		}
	}

	// Copy existing data to reduce risk of corrupting state
	if _, err := os.Stat(path); !os.IsNotExist(err) {
		d, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}
		err = ioutil.WriteFile(path+PreviousStateFilePostfix, d, DefaultStateMode)
		if err != nil {
			return err
		}
	}

	// Finally write the state
	d, err := json.MarshalIndent(s, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(path, d, DefaultStateMode)
}

// Sources of TTUs
var secretTTUSources = []func(*SecretState) (*time.Time, error){
	ttuFromTTLOrLeaseDuration,
	ttuFromCertificateValidity,
}

func ttuFromTTLOrLeaseDuration(s *SecretState) (*time.Time, error) {
	ttl, ttlKnown := s.TTL()

	var duration int
	switch {
	case ttlKnown && s.LeaseDuration > 0:
		if ttl < s.LeaseDuration {
			duration = ttl
		} else {
			duration = s.LeaseDuration
		}
	case ttlKnown:
		duration = ttl
	case s.LeaseDuration > 0:
		duration = s.LeaseDuration
	default:
		return nil, nil
	}

	ttu := s.Timestamp.Add(time.Duration(float64(duration)*s.Ratio()) * time.Second)
	return &ttu, nil
}

func ttuFromCertificateValidity(s *SecretState) (*time.Time, error) {
	if s.Data == nil {
		return nil, nil
	}

	data, ok := s.Data["certificate"].(string)
	if !ok {
		return nil, nil
	}

	block, _ := pem.Decode([]byte(data))
	if block == nil {
		return nil, fmt.Errorf("failed to parse certificate PEM")
	}

	certificate, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return nil, err
	}

	ttl := certificate.NotAfter.Sub(certificate.NotBefore)
	ttu := certificate.NotBefore.Add(time.Duration(float64(ttl) * s.Ratio()))
	return &ttu, nil
}

func (s *PouchState) SetSecret(name string, secret *api.Secret) {
	if s.Secrets == nil {
		s.Secrets = make(map[string]*SecretState)
	}
	state := &SecretState{
		Name:          name,
		Timestamp:     time.Now(),
		LeaseDuration: secret.LeaseDuration,
		Data:          secret.Data,
	}

	if _, known := state.TimeToUpdate(); !known {
		// Without a known TTU, we don't know when to update
		state.DisableAutoUpdate = true
	}

	if oldState, found := s.Secrets[name]; found {
		state.FilesUsing = oldState.FilesUsing
	}
	s.Secrets[name] = state
}

func (s *PouchState) DeleteSecret(name string) {
	delete(s.Secrets, name)
}

func (s *PouchState) NextUpdate() (secret *SecretState, minTTU time.Time) {
	for name := range s.Secrets {
		if s.Secrets[name].DisableAutoUpdate {
			continue
		}
		ttu, known := s.Secrets[name].TimeToUpdate()
		if known && (secret == nil || ttu.Before(minTTU)) {
			secret = s.Secrets[name]
			minTTU = ttu
		}
	}
	return
}

type PriorityFileSortedList []PriorityFile

type PriorityFile struct {
	Priority int    `json:"-"`
	Path     string `json:"path,omitempty"`
}

func (pf *PriorityFile) MarshalJSON() ([]byte, error) {
	return json.Marshal(pf.Path)
}

func (s *PriorityFileSortedList) UnmarshalJSON(data []byte) error {
	var priorityFiles []string

	if err := json.Unmarshal(data, &priorityFiles); err != nil {
		return err
	}

	// To keep the same state as when it was written, each file is assigned a
	// priority according to the order that they appear in the state file,
	for i, pf := range priorityFiles {
		*s = append(*s, PriorityFile{Path: pf, Priority: i * 10})
	}

	return nil
}

func (p PriorityFileSortedList) Len() int      { return len(p) }
func (p PriorityFileSortedList) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p PriorityFileSortedList) Less(i, j int) bool {
	if p[i].Priority != p[j].Priority {
		return p[i].Priority < p[j].Priority
	}
	return p[i].Path < p[j].Path
}

type SecretData map[string]interface{}

type SecretState struct {
	// Secret name
	Name string `json:"name,omitempty"`

	// Time when the secret was read
	Timestamp time.Time `json:"creation_time,omitempty"`

	// Lease duration, in seconds, if any when the secret was read
	LeaseDuration int `json:"lease_duration,omitempty"`

	// Secret will be renewed after this portion of its life has passed
	DurationRatio float64 `json:"duration_ratio,omitempty"`

	// If the secret has no expiration data, don't try to update it
	DisableAutoUpdate bool `json:"disable_auto_uptdate,omitempty"`

	// Actual secret
	Data SecretData `json:"data,omitempty"`

	// Files using this secret
	FilesUsing PriorityFileSortedList `json:"files_using,omitempty"`
}

func (s *SecretState) Ratio() float64 {
	ratio := s.DurationRatio
	if ratio == 0 {
		ratio = DefaultSecretDurationRatio
	}
	return ratio
}

func (s *SecretState) TTL() (int, bool) {
	if s.Data == nil {
		return 0, false
	}
	switch ttlNumber := s.Data["ttl"].(type) {
	case json.Number:
		ttl, err := ttlNumber.Int64()
		if err != nil {
			return 0, false
		}
		return int(ttl), true
	case int:
		return ttlNumber, true
	case int64:
		return int(ttlNumber), true
	}
	return 0, false
}

func (s *SecretState) TimeToUpdate() (minTTU time.Time, known bool) {
	for _, source := range secretTTUSources {
		ttu, err := source(s)
		if err != nil {
			log.Printf("Error trying to obtain TTU for secret '%s': %s", s.Name, err)
			continue
		}
		if ttu != nil && (!known || ttu.Before(minTTU)) {
			minTTU = *ttu
			known = true
		}
	}
	return
}

func (s *SecretState) RegisterUsage(path string, priority int) {
	for _, f := range s.FilesUsing {
		if f.Path == path {
			// Already registered
			return
		}
	}
	s.FilesUsing = append(s.FilesUsing, PriorityFile{Priority: priority, Path: path})
	sort.Sort(s.FilesUsing)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net"
	"os"
	"sync"
	"time"

	"github.com/cyberdelia/go-metrics-graphite"
	"github.com/rcrowley/go-metrics"
)

type Stats struct {
	registry metrics.Registry
	lock     sync.Mutex
}

func (t *Stats) Gauge(name string, value int) {
	metrics.GetOrRegisterGauge(name, t.registry).Update(int64(value))
}

func (t *Stats) Inc(stat string) {
	metrics.GetOrRegisterCounter(stat, t.registry).Inc(1)
}

func (t *Stats) GetCount(stat string) int64 {
	c := t.registry.Get(stat)
	if c == nil {
		return 0
	}
	return c.(metrics.Counter).Count()
}

func (t *Stats) Timing(stat string, d time.Duration) {
	metrics.GetOrRegisterTimer(stat, t.registry).Update(d)
}

func NewStats(sendToConsole bool, sendToGraphite, graphitePrefix string) *Stats {
	s := new(Stats)
	s.registry = metrics.NewRegistry()

	if sendToGraphite != "" {
		log.Println("Stats reporting to graphite: ", sendToGraphite)
		addr, _ := net.ResolveTCPAddr("tcp", sendToGraphite)

		go s.graphiteSender(time.Second*5, addr, graphitePrefix)
	}

	if sendToConsole {
		log.Println("Stats reporting enabled...")
		go metrics.Log(s.registry, time.Minute, log.New(os.Stderr, "metrics: ", log.Lmicroseconds))
	}
	return s
}

func (s *Stats) graphiteSender(freq time.Duration, addr *net.TCPAddr, prefix string) {
	cfg := graphite.GraphiteConfig{
		Addr:          addr,
		Registry:      s.registry,
		FlushInterval: 0,
		DurationUnit:  time.Millisecond,
		Prefix:        prefix,
		Percentiles:   []float64{0.5, 0.75, 0.95, 0.99, 0.999},
	}

	for _ = range time.Tick(freq) {
		graphite.GraphiteOnce(cfg)
		s.registry.UnregisterAll() // Hack to get ostrich-like flush to 0.
	}
}
<commit_msg>longer period for more accurate stats<commit_after>package main

import (
	"log"
	"net"
	"os"
	"sync"
	"time"

	"github.com/cyberdelia/go-metrics-graphite"
	"github.com/rcrowley/go-metrics"
)

type Stats struct {
	registry metrics.Registry
	lock     sync.Mutex
}

func (t *Stats) Gauge(name string, value int) {
	metrics.GetOrRegisterGauge(name, t.registry).Update(int64(value))
}

func (t *Stats) Inc(stat string) {
	metrics.GetOrRegisterCounter(stat, t.registry).Inc(1)
}

func (t *Stats) GetCount(stat string) int64 {
	c := t.registry.Get(stat)
	if c == nil {
		return 0
	}
	return c.(metrics.Counter).Count()
}

func (t *Stats) Timing(stat string, d time.Duration) {
	metrics.GetOrRegisterTimer(stat, t.registry).Update(d)
}

func NewStats(sendToConsole bool, sendToGraphite, graphitePrefix string) *Stats {
	s := new(Stats)
	s.registry = metrics.NewRegistry()

	if sendToGraphite != "" {
		log.Println("Stats reporting to graphite: ", sendToGraphite)
		addr, _ := net.ResolveTCPAddr("tcp", sendToGraphite)

		go s.graphiteSender(time.Second*15, addr, graphitePrefix)
	}

	if sendToConsole {
		log.Println("Stats reporting enabled...")
		go metrics.Log(s.registry, time.Minute, log.New(os.Stderr, "metrics: ", log.Lmicroseconds))
	}
	return s
}

func (s *Stats) graphiteSender(freq time.Duration, addr *net.TCPAddr, prefix string) {
	cfg := graphite.GraphiteConfig{
		Addr:          addr,
		Registry:      s.registry,
		FlushInterval: 0,
		DurationUnit:  time.Millisecond,
		Prefix:        prefix,
		Percentiles:   []float64{0.5, 0.75, 0.95, 0.99, 0.999},
	}

	for _ = range time.Tick(freq) {
		graphite.GraphiteOnce(cfg)
		s.registry.UnregisterAll() // Hack to get ostrich-like flush to 0.
	}
}
<|endoftext|>