<commit_before>package main

import (
	"flag"
	"net"
	"os"

	"junta/paxos"
	"junta/proc"
	"junta/store"
	"junta/util"
	"junta/client"
	"junta/server"
)

const (
	alpha = 50
	idBits = 160
)

// Flags
var (
	listenAddr *string = flag.String("l", "", "The address to bind to. Must correspond to a single public interface.")
	attachAddr *string = flag.String("a", "", "The address of another node to attach to.")
)

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

func main() {
	flag.Parse()

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

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

	outs := make(paxos.ChanPutCloserTo)

	self := util.RandHexString(idBits)
	st := store.New()
	seqn := uint64(0)
	if *attachAddr == "" { // we are the only node in a new cluster
		seqn = addMember(st, seqn + 1, self, *listenAddr)
		seqn = claimSlot(st, seqn + 1, "1", self)
		seqn = claimLeader(st, seqn + 1, self)
	} else {
		c, err := client.Dial(*attachAddr)
		if err != nil {
			panic(err)
		}

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

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

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

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

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

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

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

	go func() {
		panic(proc.Monitor(self, st))
	}()

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

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

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

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

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

func claimLeader(st *store.Store, seqn uint64, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/j/junta/leader", self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Apply(seqn, mx)
	return seqn
}
<commit_msg>publish an address for clients to use<commit_after>package main

import (
	"flag"
	"net"
	"os"

	"junta/paxos"
	"junta/proc"
	"junta/store"
	"junta/util"
	"junta/client"
	"junta/server"
)

const (
	alpha = 50
	idBits = 160
)

// Flags
var (
	listenAddr *string = flag.String("l", "", "The address to bind to. Must correspond to a single public interface.")
	publishAddr *string = flag.String("p", "", "Address to publish in junta for client connections.")
	attachAddr *string = flag.String("a", "", "The address of another node to attach to.")
)

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

func main() {
	flag.Parse()

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

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

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

	outs := make(paxos.ChanPutCloserTo)

	self := util.RandHexString(idBits)
	st := store.New()
	seqn := uint64(0)
	if *attachAddr == "" { // we are the only node in a new cluster
		seqn = addPublicAddr(st, seqn + 1, self, *publishAddr)
		seqn = addMember(st, seqn + 1, self, *listenAddr)
		seqn = claimSlot(st, seqn + 1, "1", self)
		seqn = claimLeader(st, seqn + 1, self)
	} else {
		c, err := client.Dial(*attachAddr)
		if err != nil {
			panic(err)
		}

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

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

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

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

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

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

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

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

	go func() {
		panic(proc.Monitor(self, st))
	}()

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

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

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

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

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

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

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


<commit_before>package main

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

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

	"github.com/cbroglie/mustache"
)

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

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

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

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

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

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

func parseDataFromFile(filePath string) (interface{}, error) {
	b, err := ioutil.ReadFile(filePath)
	if err != nil {
		return nil, err
	}
	var data interface{}
	if err := yaml.Unmarshal(b, &data); err != nil {
		return nil, err
	}
	return data, nil
}
<commit_msg>Print errors to stderr<commit_after>package main

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

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

	"github.com/cbroglie/mustache"
)

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

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

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

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

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

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

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


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

package main

import (
	"context"
	"flag"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"golang.org/x/sync/errgroup"

	"github.com/coreos/prometheus-operator/pkg/alertmanager"
	"github.com/coreos/prometheus-operator/pkg/analytics"
	"github.com/coreos/prometheus-operator/pkg/api"
	prometheuscontroller "github.com/coreos/prometheus-operator/pkg/prometheus"
	"github.com/go-kit/kit/log"
)

var (
	cfg              prometheuscontroller.Config
	analyticsEnabled bool
)

func init() {
	flagset := flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	flagset.StringVar(&cfg.Host, "apiserver", "", "API Server addr, e.g. ' - NOT RECOMMENDED FOR PRODUCTION - http://127.0.0.1:8080'. Omit parameter to run in on-cluster mode and utilize the service account token.")
	flagset.StringVar(&cfg.TLSConfig.CertFile, "cert-file", "", " - NOT RECOMMENDED FOR PRODUCTION - Path to public TLS certificate file.")
	flagset.StringVar(&cfg.TLSConfig.KeyFile, "key-file", "", "- NOT RECOMMENDED FOR PRODUCTION - Path to private TLS certificate file.")
	flagset.StringVar(&cfg.TLSConfig.CAFile, "ca-file", "", "- NOT RECOMMENDED FOR PRODUCTION - Path to TLS CA file.")
	flagset.StringVar(&cfg.KubeletObject, "kubelet-service", "", "Service/Endpoints object to write kubelets into in format \"namespace/name\"")
	flagset.BoolVar(&cfg.TLSInsecure, "tls-insecure", false, "- NOT RECOMMENDED FOR PRODUCTION - Don't verify API server's CA certificate.")
	flagset.BoolVar(&analyticsEnabled, "analytics", true, "Send analytical event (Cluster Created/Deleted etc.) to Google Analytics")
	flagset.StringVar(&cfg.PrometheusConfigReloader, "prometheus-config-reloader", "quay.io/coreos/prometheus-config-reloader:v0.0.2", "Config and rule reload image")
	flagset.StringVar(&cfg.ConfigReloaderImage, "config-reloader-image", "quay.io/coreos/configmap-reload:v0.0.1", "Reload Image")
	flagset.StringVar(&cfg.AlertmanagerDefaultBaseImage, "alertmanager-default-base-image", "quay.io/prometheus/alertmanager", "Alertmanager default base image")
	flagset.StringVar(&cfg.PrometheusDefaultBaseImage, "prometheus-default-base-image", "quay.io/prometheus/prometheus", "Prometheus default base image")

	flagset.Parse(os.Args[1:])
}

func Main() int {
	logger := log.NewContext(log.NewLogfmtLogger(os.Stdout)).
		With("ts", log.DefaultTimestampUTC, "caller", log.DefaultCaller)

	r := prometheus.NewRegistry()
	r.MustRegister(prometheus.NewGoCollector())

	if analyticsEnabled {
		analytics.Enable()
	}

	po, err := prometheuscontroller.New(cfg, logger.With("component", "prometheusoperator"))
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return 1
	}

	ao, err := alertmanager.New(cfg, logger.With("component", "alertmanageroperator"))
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return 1
	}

	mux := http.NewServeMux()
	web, err := api.New(cfg, logger.With("component", "api"))
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return 1
	}

	web.Register(mux)
	l, err := net.Listen("tcp", ":8080")
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return 1
	}

	po.RegisterMetrics(r)
	ao.RegisterMetrics(r)
	mux.Handle("/metrics", promhttp.HandlerFor(r, promhttp.HandlerOpts{}))

	ctx, cancel := context.WithCancel(context.Background())
	wg, ctx := errgroup.WithContext(ctx)

	wg.Go(func() error { return po.Run(ctx.Done()) })
	wg.Go(func() error { return ao.Run(ctx.Done()) })

	srv := &http.Server{Handler: mux}
	go srv.Serve(l)

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)

	select {
	case <-term:
		logger.Log("msg", "Received SIGTERM, exiting gracefully...")
	case <-ctx.Done():
	}

	cancel()
	if err := wg.Wait(); err != nil {
		logger.Log("msg", "Unhandled error received. Exiting...", "err", err)
		return 1
	}

	return 0
}

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

package main

import (
	"context"
	"flag"
	"fmt"
	"net"
	"net/http"
	"net/http/pprof"
	"os"
	"os/signal"
	"syscall"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"golang.org/x/sync/errgroup"

	"github.com/coreos/prometheus-operator/pkg/alertmanager"
	"github.com/coreos/prometheus-operator/pkg/analytics"
	"github.com/coreos/prometheus-operator/pkg/api"
	prometheuscontroller "github.com/coreos/prometheus-operator/pkg/prometheus"
	"github.com/go-kit/kit/log"
)

var (
	cfg              prometheuscontroller.Config
	analyticsEnabled bool
)

func init() {
	flagset := flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	flagset.StringVar(&cfg.Host, "apiserver", "", "API Server addr, e.g. ' - NOT RECOMMENDED FOR PRODUCTION - http://127.0.0.1:8080'. Omit parameter to run in on-cluster mode and utilize the service account token.")
	flagset.StringVar(&cfg.TLSConfig.CertFile, "cert-file", "", " - NOT RECOMMENDED FOR PRODUCTION - Path to public TLS certificate file.")
	flagset.StringVar(&cfg.TLSConfig.KeyFile, "key-file", "", "- NOT RECOMMENDED FOR PRODUCTION - Path to private TLS certificate file.")
	flagset.StringVar(&cfg.TLSConfig.CAFile, "ca-file", "", "- NOT RECOMMENDED FOR PRODUCTION - Path to TLS CA file.")
	flagset.StringVar(&cfg.KubeletObject, "kubelet-service", "", "Service/Endpoints object to write kubelets into in format \"namespace/name\"")
	flagset.BoolVar(&cfg.TLSInsecure, "tls-insecure", false, "- NOT RECOMMENDED FOR PRODUCTION - Don't verify API server's CA certificate.")
	flagset.BoolVar(&analyticsEnabled, "analytics", true, "Send analytical event (Cluster Created/Deleted etc.) to Google Analytics")
	flagset.StringVar(&cfg.PrometheusConfigReloader, "prometheus-config-reloader", "quay.io/coreos/prometheus-config-reloader:v0.0.2", "Config and rule reload image")
	flagset.StringVar(&cfg.ConfigReloaderImage, "config-reloader-image", "quay.io/coreos/configmap-reload:v0.0.1", "Reload Image")
	flagset.StringVar(&cfg.AlertmanagerDefaultBaseImage, "alertmanager-default-base-image", "quay.io/prometheus/alertmanager", "Alertmanager default base image")
	flagset.StringVar(&cfg.PrometheusDefaultBaseImage, "prometheus-default-base-image", "quay.io/prometheus/prometheus", "Prometheus default base image")

	flagset.Parse(os.Args[1:])
}

func Main() int {
	logger := log.NewContext(log.NewLogfmtLogger(os.Stdout)).
		With("ts", log.DefaultTimestampUTC, "caller", log.DefaultCaller)

	r := prometheus.NewRegistry()
	r.MustRegister(prometheus.NewGoCollector())

	if analyticsEnabled {
		analytics.Enable()
	}

	po, err := prometheuscontroller.New(cfg, logger.With("component", "prometheusoperator"))
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return 1
	}

	ao, err := alertmanager.New(cfg, logger.With("component", "alertmanageroperator"))
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return 1
	}

	mux := http.NewServeMux()
	web, err := api.New(cfg, logger.With("component", "api"))
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return 1
	}

	web.Register(mux)
	l, err := net.Listen("tcp", ":8080")
	if err != nil {
		fmt.Fprint(os.Stderr, err)
		return 1
	}

	po.RegisterMetrics(r)
	ao.RegisterMetrics(r)
	mux.Handle("/metrics", promhttp.HandlerFor(r, promhttp.HandlerOpts{}))

	mux.Handle("/debug/pprof/", http.HandlerFunc(pprof.Index))
	mux.Handle("/debug/pprof/cmdline", http.HandlerFunc(pprof.Cmdline))
	mux.Handle("/debug/pprof/profile", http.HandlerFunc(pprof.Profile))
	mux.Handle("/debug/pprof/symbol", http.HandlerFunc(pprof.Symbol))
	mux.Handle("/debug/pprof/trace", http.HandlerFunc(pprof.Trace))

	ctx, cancel := context.WithCancel(context.Background())
	wg, ctx := errgroup.WithContext(ctx)

	wg.Go(func() error { return po.Run(ctx.Done()) })
	wg.Go(func() error { return ao.Run(ctx.Done()) })

	srv := &http.Server{Handler: mux}
	go srv.Serve(l)

	term := make(chan os.Signal)
	signal.Notify(term, os.Interrupt, syscall.SIGTERM)

	select {
	case <-term:
		logger.Log("msg", "Received SIGTERM, exiting gracefully...")
	case <-ctx.Done():
	}

	cancel()
	if err := wg.Wait(); err != nil {
		logger.Log("msg", "Unhandled error received. Exiting...", "err", err)
		return 1
	}

	return 0
}

func main() {
	os.Exit(Main())
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Rook Authors. 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 ceph

import (
	"encoding/json"
	"os"
	"strings"

	"k8s.io/client-go/kubernetes"

	"github.com/pkg/errors"
	"github.com/rook/rook/cmd/rook/rook"
	osddaemon "github.com/rook/rook/pkg/daemon/ceph/osd"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mon"
	oposd "github.com/rook/rook/pkg/operator/ceph/cluster/osd"
	osdcfg "github.com/rook/rook/pkg/operator/ceph/cluster/osd/config"
	opcontroller "github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/k8sutil"
	"github.com/rook/rook/pkg/util/flags"
	"github.com/spf13/cobra"
)

var osdCmd = &cobra.Command{
	Use:   "osd",
	Short: "Provisions and runs the osd daemon",
}
var osdConfigCmd = &cobra.Command{
	Use:   "init",
	Short: "Updates ceph.conf for the osd",
}
var provisionCmd = &cobra.Command{
	Use:   "provision",
	Short: "Generates osd config and prepares an osd for runtime",
}
var osdStartCmd = &cobra.Command{
	Use:   "start",
	Short: "Starts the osd daemon", // OSDs that were provisioned by ceph-volume
}
var osdRemoveCmd = &cobra.Command{
	Use:   "remove",
	Short: "Removes a set of OSDs from the cluster",
}

var (
	osdDataDeviceFilter     string
	osdDataDevicePathFilter string
	ownerRefID              string
	mountSourcePath         string
	mountPath               string
	osdID                   int
	copyBinariesPath        string
	osdStoreType            string
	osdStringID             string
	osdUUID                 string
	osdIsDevice             bool
	pvcBackedOSD            bool
	blockPath               string
	lvBackedPV              bool
	driveGroups             string
	osdIDsToRemove          string
)

func addOSDFlags(command *cobra.Command) {
	addOSDConfigFlags(osdConfigCmd)
	addOSDConfigFlags(provisionCmd)

	// flags specific to provisioning
	provisionCmd.Flags().StringVar(&driveGroups, "drive-groups", "", "JSON marshalled specification of Ceph Drive Groups")
	provisionCmd.Flags().StringVar(&cfg.devices, "data-devices", "", "comma separated list of devices to use for storage")
	provisionCmd.Flags().StringVar(&osdDataDeviceFilter, "data-device-filter", "", "a regex filter for the device names to use, or \"all\"")
	provisionCmd.Flags().StringVar(&osdDataDevicePathFilter, "data-device-path-filter", "", "a regex filter for the device path names to use")
	provisionCmd.Flags().StringVar(&cfg.metadataDevice, "metadata-device", "", "device to use for metadata (e.g. a high performance SSD/NVMe device)")
	provisionCmd.Flags().BoolVar(&cfg.forceFormat, "force-format", false,
		"true to force the format of any specified devices, even if they already have a filesystem.  BE CAREFUL!")
	provisionCmd.Flags().BoolVar(&cfg.pvcBacked, "pvc-backed-osd", false, "true to specify a block mode pvc is backing the OSD")
	// flags for generating the osd config
	osdConfigCmd.Flags().IntVar(&osdID, "osd-id", -1, "osd id for which to generate config")
	osdConfigCmd.Flags().BoolVar(&osdIsDevice, "is-device", false, "whether the osd is a device")

	// flags for running osds that were provisioned by ceph-volume
	osdStartCmd.Flags().StringVar(&osdStringID, "osd-id", "", "the osd ID")
	osdStartCmd.Flags().StringVar(&osdUUID, "osd-uuid", "", "the osd UUID")
	osdStartCmd.Flags().StringVar(&osdStoreType, "osd-store-type", "", "whether the osd is bluestore or filestore")
	osdStartCmd.Flags().BoolVar(&pvcBackedOSD, "pvc-backed-osd", false, "Whether the OSD backing store in PVC or not")
	osdStartCmd.Flags().StringVar(&blockPath, "block-path", "", "Block path for the OSD created by ceph-volume")
	osdStartCmd.Flags().BoolVar(&lvBackedPV, "lv-backed-pv", false, "Whether the PV located on LV")

	// flags for removing OSDs that are unhealthy or otherwise should be purged from the cluster
	osdRemoveCmd.Flags().StringVar(&osdIDsToRemove, "osd-ids", "", "OSD IDs to remove from the cluster")

	// add the subcommands to the parent osd command
	osdCmd.AddCommand(osdConfigCmd,
		provisionCmd,
		osdStartCmd,
		osdRemoveCmd)
}

func addOSDConfigFlags(command *cobra.Command) {
	command.Flags().StringVar(&ownerRefID, "cluster-id", "", "the UID of the cluster CRD that owns this cluster")
	command.Flags().StringVar(&cfg.location, "location", "", "location of this node for CRUSH placement")
	command.Flags().StringVar(&cfg.nodeName, "node-name", os.Getenv("HOSTNAME"), "the host name of the node")

	// OSD store config flags
	command.Flags().IntVar(&cfg.storeConfig.WalSizeMB, "osd-wal-size", osdcfg.WalDefaultSizeMB, "default size (MB) for OSD write ahead log (WAL) (bluestore)")
	command.Flags().IntVar(&cfg.storeConfig.DatabaseSizeMB, "osd-database-size", 0, "default size (MB) for OSD database (bluestore)")
	command.Flags().StringVar(&cfg.storeConfig.StoreType, "osd-store", "", "type of backing OSD store to use (bluestore or filestore)")
	command.Flags().IntVar(&cfg.storeConfig.OSDsPerDevice, "osds-per-device", 1, "the number of OSDs per device")
	command.Flags().BoolVar(&cfg.storeConfig.EncryptedDevice, "encrypted-device", false, "whether to encrypt the OSD with dmcrypt")
}

func init() {
	addOSDFlags(osdCmd)
	addCephFlags(osdCmd)
	flags.SetFlagsFromEnv(osdCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(osdConfigCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(provisionCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(osdStartCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(osdRemoveCmd.Flags(), rook.RookEnvVarPrefix)

	osdConfigCmd.RunE = writeOSDConfig
	provisionCmd.RunE = prepareOSD
	osdStartCmd.RunE = startOSD
	osdRemoveCmd.RunE = removeOSDs
}

// Start the osd daemon if provisioned by ceph-volume
func startOSD(cmd *cobra.Command, args []string) error {
	required := []string{"osd-id", "osd-uuid"}
	if err := flags.VerifyRequiredFlags(osdStartCmd, required); err != nil {
		return err
	}

	commonOSDInit(osdStartCmd)

	context := createContext()

	// Run OSD start sequence
	err := osddaemon.StartOSD(context, osdStoreType, osdStringID, osdUUID, blockPath, pvcBackedOSD, lvBackedPV, args)
	if err != nil {
		rook.TerminateFatal(err)
	}
	return nil
}

func verifyConfigFlags(configCmd *cobra.Command) error {
	required := []string{"cluster-id", "node-name"}
	if err := flags.VerifyRequiredFlags(configCmd, required); err != nil {
		return err
	}
	required = []string{"mon-endpoints", "mon-secret", "ceph-username", "ceph-secret"}
	if err := flags.VerifyRequiredFlags(osdCmd, required); err != nil {
		return err
	}
	return nil
}

func writeOSDConfig(cmd *cobra.Command, args []string) error {
	if err := verifyConfigFlags(osdConfigCmd); err != nil {
		return err
	}
	if osdID == -1 {
		return errors.New("osd id not specified")
	}

	commonOSDInit(osdConfigCmd)

	return nil
}

// Provision a device or directory for an OSD
func prepareOSD(cmd *cobra.Command, args []string) error {
	if err := verifyConfigFlags(provisionCmd); err != nil {
		return err
	}

	var dgs osdcfg.DriveGroupBlobs
	if driveGroups != "" {
		err := json.Unmarshal([]byte(driveGroups), &dgs)
		if err != nil {
			return errors.Wrap(err, "failed to unmarshal Ceph Drive Groups spec")
		}
	}

	var dataDevices []osddaemon.DesiredDevice
	if osdDataDeviceFilter != "" {
		if cfg.devices != "" || osdDataDevicePathFilter != "" {
			return errors.New("only one of --data-devices, --data-device-filter and --data-device-path-filter can be specified")
		}

		dataDevices = []osddaemon.DesiredDevice{
			{Name: osdDataDeviceFilter, IsFilter: true, OSDsPerDevice: cfg.storeConfig.OSDsPerDevice},
		}
	} else if osdDataDevicePathFilter != "" {
		if cfg.devices != "" {
			return errors.New("only one of --data-devices, --data-device-filter and --data-device-path-filter can be specified")
		}

		dataDevices = []osddaemon.DesiredDevice{
			{Name: osdDataDevicePathFilter, IsDevicePathFilter: true, OSDsPerDevice: cfg.storeConfig.OSDsPerDevice},
		}
	} else {
		var err error
		dataDevices, err = parseDevices(cfg.devices)
		if err != nil {
			rook.TerminateFatal(errors.Wrapf(err, "failed to parse device list (%q)", cfg.devices))
		}
	}

	context := createContext()
	commonOSDInit(provisionCmd)
	crushLocation, err := getLocation(context.Clientset)
	if err != nil {
		rook.TerminateFatal(err)
	}
	logger.Infof("crush location of osd: %s", crushLocation)

	forceFormat := false
	ownerRef := opcontroller.ClusterOwnerRef(clusterInfo.Namespace, ownerRefID)
	clusterInfo.OwnerRef = ownerRef
	kv := k8sutil.NewConfigMapKVStore(clusterInfo.Namespace, context.Clientset, ownerRef)
	agent := osddaemon.NewAgent(context, dgs, dataDevices, cfg.metadataDevice, forceFormat,
		cfg.storeConfig, &clusterInfo, cfg.nodeName, kv, cfg.pvcBacked)

	err = osddaemon.Provision(context, agent, crushLocation)
	if err != nil {
		// something failed in the OSD orchestration, update the status map with failure details
		status := oposd.OrchestrationStatus{
			Status:       oposd.OrchestrationStatusFailed,
			Message:      err.Error(),
			PvcBackedOSD: cfg.pvcBacked,
		}
		oposd.UpdateNodeStatus(kv, cfg.nodeName, status)

		rook.TerminateFatal(err)
	}

	return nil
}

// Purge the desired OSDs from the cluster
func removeOSDs(cmd *cobra.Command, args []string) error {
	required := []string{"osd-ids"}
	if err := flags.VerifyRequiredFlags(osdRemoveCmd, required); err != nil {
		return err
	}
	required = []string{"mon-endpoints", "ceph-username", "ceph-secret"}
	if err := flags.VerifyRequiredFlags(osdCmd, required); err != nil {
		return err
	}

	commonOSDInit(osdRemoveCmd)

	context := createContext()

	// Run OSD remove sequence
	err := osddaemon.RemoveOSDs(context, &clusterInfo, strings.Split(osdIDsToRemove, ","))
	if err != nil {
		rook.TerminateFatal(err)
	}
	return nil
}

func commonOSDInit(cmd *cobra.Command) {
	rook.SetLogLevel()
	rook.LogStartupInfo(cmd.Flags())

	clusterInfo.Monitors = mon.ParseMonEndpoints(cfg.monEndpoints)
}

// use zone/region/hostname labels in the crushmap
func getLocation(clientset kubernetes.Interface) (string, error) {
	// get the value the operator instructed to use as the host name in the CRUSH map
	hostNameLabel := os.Getenv("ROOK_CRUSHMAP_HOSTNAME")

	loc, err := oposd.GetLocationWithNode(clientset, os.Getenv(k8sutil.NodeNameEnvVar), hostNameLabel)
	if err != nil {
		return "", err
	}
	return loc, nil
}

func updateLocationWithNodeLabels(location *[]string, nodeLabels map[string]string) {
	oposd.UpdateLocationWithNodeLabels(location, nodeLabels)
}

// Parse the devices, which are sent as a JSON-marshalled list of device IDs with a StorageConfig spec
func parseDevices(devices string) ([]osddaemon.DesiredDevice, error) {
	configuredDevices := []osdcfg.ConfiguredDevice{}
	err := json.Unmarshal([]byte(devices), &configuredDevices)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to JSON unmarshal configured devices (%q)", devices)
	}

	var result []osddaemon.DesiredDevice
	for _, cd := range configuredDevices {
		d := osddaemon.DesiredDevice{
			Name: cd.ID,
		}
		d.OSDsPerDevice = cd.StoreConfig.OSDsPerDevice
		d.DatabaseSizeMB = cd.StoreConfig.DatabaseSizeMB
		d.DeviceClass = cd.StoreConfig.DeviceClass
		d.MetadataDevice = cd.StoreConfig.MetadataDevice

		if d.OSDsPerDevice < 1 {
			return nil, errors.Errorf("osds per device should be greater than 0 (%q)", d.OSDsPerDevice)
		}

		result = append(result, d)
	}

	logger.Infof("desired devices to configure osds: %+v", result)
	return result, nil
}
<commit_msg>ceph: fix drive group deployment failure<commit_after>/*
Copyright 2016 The Rook Authors. 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 ceph

import (
	"encoding/json"
	"os"
	"strings"

	"k8s.io/client-go/kubernetes"

	"github.com/pkg/errors"
	"github.com/rook/rook/cmd/rook/rook"
	osddaemon "github.com/rook/rook/pkg/daemon/ceph/osd"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mon"
	oposd "github.com/rook/rook/pkg/operator/ceph/cluster/osd"
	osdcfg "github.com/rook/rook/pkg/operator/ceph/cluster/osd/config"
	opcontroller "github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/k8sutil"
	"github.com/rook/rook/pkg/util/flags"
	"github.com/spf13/cobra"
)

var osdCmd = &cobra.Command{
	Use:   "osd",
	Short: "Provisions and runs the osd daemon",
}
var osdConfigCmd = &cobra.Command{
	Use:   "init",
	Short: "Updates ceph.conf for the osd",
}
var provisionCmd = &cobra.Command{
	Use:   "provision",
	Short: "Generates osd config and prepares an osd for runtime",
}
var osdStartCmd = &cobra.Command{
	Use:   "start",
	Short: "Starts the osd daemon", // OSDs that were provisioned by ceph-volume
}
var osdRemoveCmd = &cobra.Command{
	Use:   "remove",
	Short: "Removes a set of OSDs from the cluster",
}

var (
	osdDataDeviceFilter     string
	osdDataDevicePathFilter string
	ownerRefID              string
	mountSourcePath         string
	mountPath               string
	osdID                   int
	copyBinariesPath        string
	osdStoreType            string
	osdStringID             string
	osdUUID                 string
	osdIsDevice             bool
	pvcBackedOSD            bool
	blockPath               string
	lvBackedPV              bool
	driveGroups             string
	osdIDsToRemove          string
)

func addOSDFlags(command *cobra.Command) {
	addOSDConfigFlags(osdConfigCmd)
	addOSDConfigFlags(provisionCmd)

	// flags specific to provisioning
	provisionCmd.Flags().StringVar(&driveGroups, "drive-groups", "", "JSON marshalled specification of Ceph Drive Groups")
	provisionCmd.Flags().StringVar(&cfg.devices, "data-devices", "", "comma separated list of devices to use for storage")
	provisionCmd.Flags().StringVar(&osdDataDeviceFilter, "data-device-filter", "", "a regex filter for the device names to use, or \"all\"")
	provisionCmd.Flags().StringVar(&osdDataDevicePathFilter, "data-device-path-filter", "", "a regex filter for the device path names to use")
	provisionCmd.Flags().StringVar(&cfg.metadataDevice, "metadata-device", "", "device to use for metadata (e.g. a high performance SSD/NVMe device)")
	provisionCmd.Flags().BoolVar(&cfg.forceFormat, "force-format", false,
		"true to force the format of any specified devices, even if they already have a filesystem.  BE CAREFUL!")
	provisionCmd.Flags().BoolVar(&cfg.pvcBacked, "pvc-backed-osd", false, "true to specify a block mode pvc is backing the OSD")
	// flags for generating the osd config
	osdConfigCmd.Flags().IntVar(&osdID, "osd-id", -1, "osd id for which to generate config")
	osdConfigCmd.Flags().BoolVar(&osdIsDevice, "is-device", false, "whether the osd is a device")

	// flags for running osds that were provisioned by ceph-volume
	osdStartCmd.Flags().StringVar(&osdStringID, "osd-id", "", "the osd ID")
	osdStartCmd.Flags().StringVar(&osdUUID, "osd-uuid", "", "the osd UUID")
	osdStartCmd.Flags().StringVar(&osdStoreType, "osd-store-type", "", "whether the osd is bluestore or filestore")
	osdStartCmd.Flags().BoolVar(&pvcBackedOSD, "pvc-backed-osd", false, "Whether the OSD backing store in PVC or not")
	osdStartCmd.Flags().StringVar(&blockPath, "block-path", "", "Block path for the OSD created by ceph-volume")
	osdStartCmd.Flags().BoolVar(&lvBackedPV, "lv-backed-pv", false, "Whether the PV located on LV")

	// flags for removing OSDs that are unhealthy or otherwise should be purged from the cluster
	osdRemoveCmd.Flags().StringVar(&osdIDsToRemove, "osd-ids", "", "OSD IDs to remove from the cluster")

	// add the subcommands to the parent osd command
	osdCmd.AddCommand(osdConfigCmd,
		provisionCmd,
		osdStartCmd,
		osdRemoveCmd)
}

func addOSDConfigFlags(command *cobra.Command) {
	command.Flags().StringVar(&ownerRefID, "cluster-id", "", "the UID of the cluster CRD that owns this cluster")
	command.Flags().StringVar(&cfg.location, "location", "", "location of this node for CRUSH placement")
	command.Flags().StringVar(&cfg.nodeName, "node-name", os.Getenv("HOSTNAME"), "the host name of the node")

	// OSD store config flags
	command.Flags().IntVar(&cfg.storeConfig.WalSizeMB, "osd-wal-size", osdcfg.WalDefaultSizeMB, "default size (MB) for OSD write ahead log (WAL) (bluestore)")
	command.Flags().IntVar(&cfg.storeConfig.DatabaseSizeMB, "osd-database-size", 0, "default size (MB) for OSD database (bluestore)")
	command.Flags().StringVar(&cfg.storeConfig.StoreType, "osd-store", "", "type of backing OSD store to use (bluestore or filestore)")
	command.Flags().IntVar(&cfg.storeConfig.OSDsPerDevice, "osds-per-device", 1, "the number of OSDs per device")
	command.Flags().BoolVar(&cfg.storeConfig.EncryptedDevice, "encrypted-device", false, "whether to encrypt the OSD with dmcrypt")
}

func init() {
	addOSDFlags(osdCmd)
	addCephFlags(osdCmd)
	flags.SetFlagsFromEnv(osdCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(osdConfigCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(provisionCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(osdStartCmd.Flags(), rook.RookEnvVarPrefix)
	flags.SetFlagsFromEnv(osdRemoveCmd.Flags(), rook.RookEnvVarPrefix)

	osdConfigCmd.RunE = writeOSDConfig
	provisionCmd.RunE = prepareOSD
	osdStartCmd.RunE = startOSD
	osdRemoveCmd.RunE = removeOSDs
}

// Start the osd daemon if provisioned by ceph-volume
func startOSD(cmd *cobra.Command, args []string) error {
	required := []string{"osd-id", "osd-uuid"}
	if err := flags.VerifyRequiredFlags(osdStartCmd, required); err != nil {
		return err
	}

	commonOSDInit(osdStartCmd)

	context := createContext()

	// Run OSD start sequence
	err := osddaemon.StartOSD(context, osdStoreType, osdStringID, osdUUID, blockPath, pvcBackedOSD, lvBackedPV, args)
	if err != nil {
		rook.TerminateFatal(err)
	}
	return nil
}

func verifyConfigFlags(configCmd *cobra.Command) error {
	required := []string{"cluster-id", "node-name"}
	if err := flags.VerifyRequiredFlags(configCmd, required); err != nil {
		return err
	}
	required = []string{"mon-endpoints", "mon-secret", "ceph-username", "ceph-secret"}
	if err := flags.VerifyRequiredFlags(osdCmd, required); err != nil {
		return err
	}
	return nil
}

func writeOSDConfig(cmd *cobra.Command, args []string) error {
	if err := verifyConfigFlags(osdConfigCmd); err != nil {
		return err
	}
	if osdID == -1 {
		return errors.New("osd id not specified")
	}

	commonOSDInit(osdConfigCmd)

	return nil
}

// Provision a device or directory for an OSD
func prepareOSD(cmd *cobra.Command, args []string) error {
	if err := verifyConfigFlags(provisionCmd); err != nil {
		return err
	}

	var dgs osdcfg.DriveGroupBlobs
	if driveGroups != "" {
		err := json.Unmarshal([]byte(driveGroups), &dgs)
		if err != nil {
			return errors.Wrap(err, "failed to unmarshal Ceph Drive Groups spec")
		}
	}

	var dataDevices []osddaemon.DesiredDevice
	if osdDataDeviceFilter != "" {
		if cfg.devices != "" || osdDataDevicePathFilter != "" {
			return errors.New("only one of --data-devices, --data-device-filter and --data-device-path-filter can be specified")
		}

		dataDevices = []osddaemon.DesiredDevice{
			{Name: osdDataDeviceFilter, IsFilter: true, OSDsPerDevice: cfg.storeConfig.OSDsPerDevice},
		}
	} else if osdDataDevicePathFilter != "" {
		if cfg.devices != "" {
			return errors.New("only one of --data-devices, --data-device-filter and --data-device-path-filter can be specified")
		}

		dataDevices = []osddaemon.DesiredDevice{
			{Name: osdDataDevicePathFilter, IsDevicePathFilter: true, OSDsPerDevice: cfg.storeConfig.OSDsPerDevice},
		}
	} else {
		var err error
		dataDevices, err = parseDevices(cfg.devices)
		if err != nil {
			rook.TerminateFatal(errors.Wrapf(err, "failed to parse device list (%q)", cfg.devices))
		}
	}

	context := createContext()
	commonOSDInit(provisionCmd)
	crushLocation, err := getLocation(context.Clientset)
	if err != nil {
		rook.TerminateFatal(err)
	}
	logger.Infof("crush location of osd: %s", crushLocation)

	forceFormat := false
	ownerRef := opcontroller.ClusterOwnerRef(clusterInfo.Namespace, ownerRefID)
	clusterInfo.OwnerRef = ownerRef
	kv := k8sutil.NewConfigMapKVStore(clusterInfo.Namespace, context.Clientset, ownerRef)
	agent := osddaemon.NewAgent(context, dgs, dataDevices, cfg.metadataDevice, forceFormat,
		cfg.storeConfig, &clusterInfo, cfg.nodeName, kv, cfg.pvcBacked)

	err = osddaemon.Provision(context, agent, crushLocation)
	if err != nil {
		// something failed in the OSD orchestration, update the status map with failure details
		status := oposd.OrchestrationStatus{
			Status:       oposd.OrchestrationStatusFailed,
			Message:      err.Error(),
			PvcBackedOSD: cfg.pvcBacked,
		}
		oposd.UpdateNodeStatus(kv, cfg.nodeName, status)

		rook.TerminateFatal(err)
	}

	return nil
}

// Purge the desired OSDs from the cluster
func removeOSDs(cmd *cobra.Command, args []string) error {
	required := []string{"osd-ids"}
	if err := flags.VerifyRequiredFlags(osdRemoveCmd, required); err != nil {
		return err
	}
	required = []string{"mon-endpoints", "ceph-username", "ceph-secret"}
	if err := flags.VerifyRequiredFlags(osdCmd, required); err != nil {
		return err
	}

	commonOSDInit(osdRemoveCmd)

	context := createContext()

	// Run OSD remove sequence
	err := osddaemon.RemoveOSDs(context, &clusterInfo, strings.Split(osdIDsToRemove, ","))
	if err != nil {
		rook.TerminateFatal(err)
	}
	return nil
}

func commonOSDInit(cmd *cobra.Command) {
	rook.SetLogLevel()
	rook.LogStartupInfo(cmd.Flags())

	clusterInfo.Monitors = mon.ParseMonEndpoints(cfg.monEndpoints)
}

// use zone/region/hostname labels in the crushmap
func getLocation(clientset kubernetes.Interface) (string, error) {
	// get the value the operator instructed to use as the host name in the CRUSH map
	hostNameLabel := os.Getenv("ROOK_CRUSHMAP_HOSTNAME")

	loc, err := oposd.GetLocationWithNode(clientset, os.Getenv(k8sutil.NodeNameEnvVar), hostNameLabel)
	if err != nil {
		return "", err
	}
	return loc, nil
}

func updateLocationWithNodeLabels(location *[]string, nodeLabels map[string]string) {
	oposd.UpdateLocationWithNodeLabels(location, nodeLabels)
}

// Parse the devices, which are sent as a JSON-marshalled list of device IDs with a StorageConfig spec
func parseDevices(devices string) ([]osddaemon.DesiredDevice, error) {
	if devices == "" {
		return []osddaemon.DesiredDevice{}, nil
	}

	configuredDevices := []osdcfg.ConfiguredDevice{}
	err := json.Unmarshal([]byte(devices), &configuredDevices)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to JSON unmarshal configured devices (%q)", devices)
	}

	var result []osddaemon.DesiredDevice
	for _, cd := range configuredDevices {
		d := osddaemon.DesiredDevice{
			Name: cd.ID,
		}
		d.OSDsPerDevice = cd.StoreConfig.OSDsPerDevice
		d.DatabaseSizeMB = cd.StoreConfig.DatabaseSizeMB
		d.DeviceClass = cd.StoreConfig.DeviceClass
		d.MetadataDevice = cd.StoreConfig.MetadataDevice

		if d.OSDsPerDevice < 1 {
			return nil, errors.Errorf("osds per device should be greater than 0 (%q)", d.OSDsPerDevice)
		}

		result = append(result, d)
	}

	logger.Infof("desired devices to configure osds: %+v", result)
	return result, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/BurntSushi/toml"
	"github.com/kaneshin/go-talentio/talentio"
)

var (
	page       = flag.Int("page", 1, "")
	status     = flag.String("status", "ongoing", "")
	configPath = flag.String("config", "$HOME/.talentio.tml", "")
)

// config represents a configuration of commands.
var config = struct {
	AccessToken string `toml:"access_token"`
}{}

func flagParse() {
	flag.Parse()

	if *page < 1 {
		*page = 1
	}

	statuses := [6]string{
		talentio.StatusOngoing,
		talentio.StatusReject,
		talentio.StatusFail,
		talentio.StatusPass,
		talentio.StatusPoolActive,
		talentio.StatusPoolInactive,
	}
	for _, s := range statuses {
		if *status == s {
			goto STATUS_OK
		}
	}
	*status = talentio.StatusOngoing
STATUS_OK:

	fp := os.ExpandEnv(*configPath)
	_, _ = toml.DecodeFile(fp, &config)
}

func run() int {
	flagParse()

	c := talentio.NewConfig().WithHTTPClient(http.DefaultClient)
	str := config.AccessToken
	if str == "" {
		str = os.Getenv("TALENTIO_ACCESS_TOKEN")
	}
	c.WithAccessToken(str)

	client := talentio.NewClient(c)

	opt := talentio.CandidatesListOptions{
		Page: *page,
		// FIXME:
		// It's not working. API might be broken.
		// Status: *status,
		Sort: talentio.SortRegisteredAtDescKey,
	}
	candidates, _, err := client.Candidates.List(&opt)
	if err != nil {
		log.Printf("error: %v", err.Error())
		return 1
	}

	for _, candidate := range candidates {
		if *status != candidate.Status {
			// FIXME:
			// The because of opt.Status is not working.
			continue
		}

		fmt.Printf("ID=%d: %v %v\n", candidate.ID, candidate.FirstName, candidate.LastName)
	}

	return 0
}

func main() {
	os.Exit(run())
}
<commit_msg>Status is working well!! Thanks official team :)<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"

	"github.com/BurntSushi/toml"
	"github.com/kaneshin/go-talentio/talentio"
)

var (
	page       = flag.Int("page", 1, "")
	status     = flag.String("status", "ongoing", "")
	configPath = flag.String("config", "$HOME/.talentio.tml", "")
)

// config represents a configuration of commands.
var config = struct {
	AccessToken string `toml:"access_token"`
}{}

func flagParse() {
	flag.Parse()

	if *page < 1 {
		*page = 1
	}

	statuses := [6]string{
		talentio.StatusOngoing,
		talentio.StatusReject,
		talentio.StatusFail,
		talentio.StatusPass,
		talentio.StatusPoolActive,
		talentio.StatusPoolInactive,
	}
	for _, s := range statuses {
		if *status == s {
			goto STATUS_OK
		}
	}
	*status = talentio.StatusOngoing
STATUS_OK:

	fp := os.ExpandEnv(*configPath)
	_, _ = toml.DecodeFile(fp, &config)
}

func run() int {
	flagParse()

	c := talentio.NewConfig().WithHTTPClient(http.DefaultClient)
	str := config.AccessToken
	if str == "" {
		str = os.Getenv("TALENTIO_ACCESS_TOKEN")
	}
	c.WithAccessToken(str)

	client := talentio.NewClient(c)

	opt := talentio.CandidatesListOptions{
		Page:   *page,
		Status: *status,
		Sort:   talentio.SortRegisteredAtDescKey,
	}
	candidates, _, err := client.Candidates.List(&opt)
	if err != nil {
		log.Printf("error: %v", err.Error())
		return 1
	}

	for _, candidate := range candidates {

		fmt.Printf("ID=%d: %v %v (%v)\n", candidate.ID, candidate.FirstName, candidate.LastName, candidate.Status)
	}

	return 0
}

func main() {
	os.Exit(run())
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var (
	output   string
	pkg      string
	function string
	fs       string

	tmpl = template.Must(template.New("go-template").Parse(goTemplate))
)

type templateData struct {
	Pkg      string
	Function string
	FS       string

	Files []File
}

type File struct {
	Path    string
	Data    string
	Size    int64
	Mode    uint32
	ModTime int64
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Fprintln(os.Stderr, "No input files/directories given")
		os.Exit(0)
	}

	if pkg == "" {
		fmt.Fprintln(os.Stderr, "No output file package given")
		os.Exit(0)
	}

	var out *os.File

	if output == "-" {
		out = os.Stdout
	} else {
		var err error
		if out, err = os.OpenFile(output, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644); err != nil {
			fmt.Fprintf(os.Stderr, "Error opening '%s' for writing: %v\n", output, err)
			os.Exit(0)
		}

		defer out.Close()
	}

	files := []File{}

	for _, arg := range flag.Args() {
		recursive := false
		if strings.HasSuffix(arg, "/...") {
			recursive = true
			arg = arg[:len(arg)-4]
		}
		func() {
			file, err := os.Open(arg)
			defer file.Close()

			if err != nil {
				fmt.Fprintf(os.Stderr, "Error opening '%s': %v\n", arg, err)
				return
			}

			stat, err := file.Stat()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error stat-ing '%s': %v\n", arg, err)
				return
			}

			if stat.IsDir() {
				filepath.Walk(arg, func(path string, info os.FileInfo, err error) error {
					if info.IsDir() {
						if !recursive {
							return filepath.SkipDir
						}
					} else {
						if b, err := ioutil.ReadFile(path); err == nil {
							files = append(files, File{
								Path:    path,
								Data:    fmt.Sprintf("%q", b),
								Size:    info.Size(),
								ModTime: info.ModTime().Unix(),
								Mode:    uint32(info.Mode()),
							})
						} else {
							fmt.Fprintf(os.Stderr, "Error reading '%s': %v\n", path, err)
						}
					}

					return nil
				})
			} else {
				if b, err := ioutil.ReadAll(file); err == nil {
					files = append(files, File{
						Path:    arg,
						Data:    fmt.Sprintf("%q", b),
						Size:    stat.Size(),
						ModTime: stat.ModTime().Unix(),
						Mode:    uint32(stat.Mode()),
					})
				} else {
					fmt.Fprintf(os.Stderr, "Error reading '%s': %v\n", arg, err)
					return
				}

			}
		}()
	}

	buf := new(bytes.Buffer)
	err := tmpl.Execute(buf, templateData{Pkg: pkg, Function: function, FS: fs, Files: files})
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing template: %v\n", err)
		os.Exit(0)
	}

	if _, err := buf.WriteTo(out); err != nil {
		fmt.Fprintf(os.Stderr, "Error writing template result to output: %v\n", err)
		os.Exit(0)
	}
}

func init() {
	flag.StringVar(&output, "output", "-", "the output file")
	flag.StringVar(&pkg, "pkg", "main", "the package of the output go file")
	flag.StringVar(&function, "function", "addFiles", "the function that will add the files to the fs")
	flag.StringVar(&fs, "fs", "DefaultFS", "The fs object to be used. Will be created if different from the default")
}

const goTemplate = `package {{ .Pkg }}

// DO NOT EDIT ** This file was generated with the webfw-fs tool ** DO NOT EDIT //

import (
	"fmt"
	"os"
	"time"

	"github.com/urandom/webfw/fs"
)

type ErrNotAdded struct {
	Path string
}

func {{ .Function }}() (*fs.FS, error) {
{{ if eq .FS "DefaultFS" }}
	wfs := fs.DefaultFS
{{ else }}
	wfs := fs.NewFS()
{{ end }}
	var (
		size int64
		mode os.FileMode
		t    time.Time
	)
{{ range $index, $file := .Files }}
	size = int64({{ $file.Size }})
	mode = os.FileMode({{ $file.Mode }})
	t = time.Unix({{ $file.ModTime }}, 0)
	if !wfs.Add(fs.NewFileListing("{{ $file.Path }}", size, mode, t, []byte({{ $file.Data }}))) {
		return nil, &ErrNotAdded{Path: "$path"}
	}
{{ end }}
	return wfs, nil
}

func (e ErrNotAdded) Error() string {
	return fmt.Sprintf("Error adding file '%s'\n", e.Path)
}
`
<commit_msg>add an option to format the generated code<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/format"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"text/template"
)

var (
	output   string
	pkg      string
	function string
	fs       string
	doFormat bool

	tmpl = template.Must(template.New("go-template").Parse(goTemplate))
)

type templateData struct {
	Pkg      string
	Function string
	FS       string

	Files []File
}

type File struct {
	Path    string
	Data    string
	Size    int64
	Mode    uint32
	ModTime int64
}

func main() {
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Fprintln(os.Stderr, "No input files/directories given")
		os.Exit(0)
	}

	if pkg == "" {
		fmt.Fprintln(os.Stderr, "No output file package given")
		os.Exit(0)
	}

	var out *os.File

	if output == "-" {
		out = os.Stdout
	} else {
		var err error
		if out, err = os.OpenFile(output, os.O_WRONLY|os.O_TRUNC|os.O_CREATE, 0644); err != nil {
			fmt.Fprintf(os.Stderr, "Error opening '%s' for writing: %v\n", output, err)
			os.Exit(0)
		}

		defer out.Close()
	}

	files := []File{}

	for _, arg := range flag.Args() {
		recursive := false
		if strings.HasSuffix(arg, "/...") {
			recursive = true
			arg = arg[:len(arg)-4]
		}
		func() {
			file, err := os.Open(arg)
			defer file.Close()

			if err != nil {
				fmt.Fprintf(os.Stderr, "Error opening '%s': %v\n", arg, err)
				return
			}

			stat, err := file.Stat()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error stat-ing '%s': %v\n", arg, err)
				return
			}

			if stat.IsDir() {
				filepath.Walk(arg, func(path string, info os.FileInfo, err error) error {
					if info.IsDir() {
						if !recursive {
							return filepath.SkipDir
						}
					} else {
						if b, err := ioutil.ReadFile(path); err == nil {
							files = append(files, File{
								Path:    path,
								Data:    fmt.Sprintf("%q", b),
								Size:    info.Size(),
								ModTime: info.ModTime().Unix(),
								Mode:    uint32(info.Mode()),
							})
						} else {
							fmt.Fprintf(os.Stderr, "Error reading '%s': %v\n", path, err)
						}
					}

					return nil
				})
			} else {
				if b, err := ioutil.ReadAll(file); err == nil {
					files = append(files, File{
						Path:    arg,
						Data:    fmt.Sprintf("%q", b),
						Size:    stat.Size(),
						ModTime: stat.ModTime().Unix(),
						Mode:    uint32(stat.Mode()),
					})
				} else {
					fmt.Fprintf(os.Stderr, "Error reading '%s': %v\n", arg, err)
					return
				}

			}
		}()
	}

	buf := new(bytes.Buffer)
	err := tmpl.Execute(buf, templateData{Pkg: pkg, Function: function, FS: fs, Files: files})
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error executing template: %v\n", err)
		os.Exit(0)
	}

	if doFormat {
		if b, err := format.Source(buf.Bytes()); err == nil {
			buf.Reset()
			if _, err = buf.Write(b); err != nil {
				fmt.Fprintf(os.Stderr, "Error writing formatted code to buffer: %v\n", err)
				os.Exit(0)
			}
		} else {
			fmt.Fprintf(os.Stderr, "Error formatting generated code: %v\n", err)
			os.Exit(0)
		}
	}

	if _, err := buf.WriteTo(out); err != nil {
		fmt.Fprintf(os.Stderr, "Error writing template result to output: %v\n", err)
		os.Exit(0)
	}
}

func init() {
	flag.StringVar(&output, "output", "-", "the output file")
	flag.StringVar(&pkg, "package", "main", "the package of the output go file")
	flag.StringVar(&function, "function", "addFiles", "the function that will add the files to the fs")
	flag.StringVar(&fs, "fs", "DefaultFS", "the fs object to be used. Will be created if different from the default")
	flag.BoolVar(&doFormat, "format", false, "run the output go code through go/format")
}

const goTemplate = `package {{ .Pkg }}

// DO NOT EDIT ** This file was generated with the webfw-fs tool ** DO NOT EDIT //

import (
	"fmt"
	"os"
	"time"

	"github.com/urandom/webfw/fs"
)

type ErrNotAdded struct {
	Path string
}

func {{ .Function }}() (*fs.FS, error) {
{{ if eq .FS "DefaultFS" }}
	wfs := fs.DefaultFS
{{ else }}
	wfs := fs.NewFS()
{{ end }}
	var (
		size int64
		mode os.FileMode
		t    time.Time
	)
{{ range $index, $file := .Files }}
	size = int64({{ $file.Size }})
	mode = os.FileMode({{ $file.Mode }})
	t = time.Unix({{ $file.ModTime }}, 0)
	if !wfs.Add(fs.NewFileListing("{{ $file.Path }}", size, mode, t, []byte({{ $file.Data }}))) {
		return nil, &ErrNotAdded{Path: "$path"}
	}
{{ end }}
	return wfs, nil
}

func (e ErrNotAdded) Error() string {
	return fmt.Sprintf("Error adding file '%s'\n", e.Path)
}
`
<|endoftext|>


<commit_before>/*
	push a zx replica
*/
package main

import (
	"clive/cmd"
	"clive/cmd/opt"
	"clive/zx/repl"
	"os"
)

var (
	opts = opt.New("file")
	notux, nflag bool
)

func main() {
	cmd.UnixIO("err")
	c := cmd.AppCtx()
	opts.NewFlag("D", "debug", &c.Debug)
	opts.NewFlag("v", "verbose", &c.Verb)
	opts.NewFlag("u", "don't use unix out", &notux)
	opts.NewFlag("n", "dry run", &nflag)
	args := opts.Parse()
	if !notux {
		cmd.UnixIO("out")
	}
	if len(args) != 1 {
		opts.Usage()
	}
	tr, err := repl.Load(args[0])
	if err != nil {
		cmd.Fatal(err)
	}
	if c.Debug {
		tr.Ldb.DumpTo(os.Stderr)
		tr.Rdb.DumpTo(os.Stderr)
	}
	if nflag {
		cc, err := tr.Changes()
		if err != nil {
			cmd.Fatal(err)
		}
		for c := range cc {
			cmd.Printf("%s\n", c)
		}
		cmd.Exit(nil)
	}
	var cc chan repl.Chg
	dc := make(chan bool)
	if c.Verb {
		cc = make(chan repl.Chg)
		go func() {
			for c := range cc {
				cmd.Printf("%s\n", c)
			}
			close(dc)
		}()
	} else {
		close(dc)
	}
	err = tr.Sync(cc)
	<-dc
	if err := tr.Save(args[0]); err != nil {
		tr.Close()
		cmd.Fatal("save: %s", err)
	}
	tr.Close()
	if err != nil {
		cmd.Fatal(err)
	}
	cmd.Exit(nil)
}
<commit_msg>zxrepl  fix<commit_after>/*
	push a zx replica
*/
package main

import (
	"clive/cmd"
	"clive/cmd/opt"
	"clive/zx/repl"
	"os"
)

var (
	opts = opt.New("file")
	notux, nflag bool
)

func main() {
	cmd.UnixIO("err")
	c := cmd.AppCtx()
	opts.NewFlag("D", "debug", &c.Debug)
	opts.NewFlag("v", "verbose", &c.Verb)
	opts.NewFlag("u", "don't use unix out", &notux)
	opts.NewFlag("n", "dry run", &nflag)
	args := opts.Parse()
	if !notux {
		cmd.UnixIO("out")
	}
	if len(args) != 1 {
		opts.Usage()
	}
	tr, err := repl.Load(args[0])
	if err != nil {
		cmd.Fatal(err)
	}
	if c.Debug {
		tr.Ldb.DumpTo(os.Stderr)
		tr.Rdb.DumpTo(os.Stderr)
	}
	if nflag {
		cc, err := tr.Changes()
		if err != nil {
			cmd.Fatal(err)
		}
		for c := range cc {
			cmd.Printf("chg %s %s\n", c.At, c)
		}
		cmd.Exit(nil)
	}
	var cc chan repl.Chg
	dc := make(chan bool)
	if c.Verb {
		cc = make(chan repl.Chg)
		go func() {
			for c := range cc {
				cmd.Printf("chg %s %s\n", c.At, c)
			}
			close(dc)
		}()
	} else {
		close(dc)
	}
	err = tr.Sync(cc)
	<-dc
	if err := tr.Save(args[0]); err != nil {
		tr.Close()
		cmd.Fatal("save: %s", err)
	}
	tr.Close()
	if err != nil {
		cmd.Fatal(err)
	}
	cmd.Exit(nil)
}
<|endoftext|>


<commit_before><commit_msg>Support GCR<commit_after><|endoftext|>


<commit_before>// Copyright (c) 2012-2020 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a MIT license found in the LICENSE file.

package codec

// This file sets up the variables used, including testInitFns.
// Each file should add initialization that should be performed
// after flags are parsed.
//
// init is a multi-step process:
//   - setup vars (handled by init functions in each file)
//   - parse flags
//   - setup derived vars (handled by pre-init registered functions - registered in init function)
//   - post init (handled by post-init registered functions - registered in init function)
// This way, no one has to manage carefully control the initialization
// using file names, etc.
//
// Tests which require external dependencies need the -tag=x parameter.
// They should be run as:
//    go test -tags=x -run=. <other parameters ...>
// Benchmarks should also take this parameter, to include the sereal, xdr, etc.
// To run against codecgen, etc, make sure you pass extra parameters.
// Example usage:
//    go test "-tags=x codecgen" -bench=. <other parameters ...>
//
// To fully test everything:
//    go test -tags=x -benchtime=100ms -tv -bg -bi  -brw -bu -v -run=. -bench=.

// Handling flags
// codec_test.go will define a set of global flags for testing, including:
//   - Use Reset
//   - Use IO reader/writer (vs direct bytes)
//   - Set Canonical
//   - Set InternStrings
//   - Use Symbols
//
// This way, we can test them all by running same set of tests with a different
// set of flags.
//
// Following this, all the benchmarks will utilize flags set by codec_test.go
// and will not redefine these "global" flags.

import (
	"bytes"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"sync"
)

import . "github.com/ugorji/go/codec"

type testHED struct {
	H Handle
	E *Encoder
	D *Decoder
}

type ioReaderWrapper struct {
	r io.Reader
}

func (x ioReaderWrapper) Read(p []byte) (n int, err error) {
	return x.r.Read(p)
}

type ioWriterWrapper struct {
	w io.Writer
}

func (x ioWriterWrapper) Write(p []byte) (n int, err error) {
	return x.w.Write(p)
}

var (
	// testNoopH    = NoopHandle(8)
	testMsgpackH = &MsgpackHandle{}
	testBincH    = &BincHandle{}
	testSimpleH  = &SimpleHandle{}
	testCborH    = &CborHandle{}
	testJsonH    = &JsonHandle{}

	testHandles     []Handle
	testPreInitFns  []func()
	testPostInitFns []func()

	testOnce sync.Once

	testHEDs []testHED
)

// flag variables used by tests (and bench)
var (
	testVerbose bool

	//depth of 0 maps to ~400bytes json-encoded string, 1 maps to ~1400 bytes, etc
	//For depth>1, we likely trigger stack growth for encoders, making benchmarking unreliable.
	testDepth int

	testMaxInitLen int

	testInitDebug bool
	testUseReset  bool
	testSkipIntf  bool
	testUseMust   bool

	testUseIoEncDec  int
	testUseIoWrapper bool

	testNumRepeatString int

	testRpcBufsize       int
	testMapStringKeyOnly bool
)

// variables that are not flags, but which can configure the handles
var (
	testEncodeOptions EncodeOptions
	testDecodeOptions DecodeOptions
)

func init() {
	log.SetOutput(ioutil.Discard) // don't allow things log to standard out/err
	testHEDs = make([]testHED, 0, 32)
	testHandles = append(testHandles,
		// testNoopH,
		testMsgpackH, testBincH, testSimpleH, testCborH, testJsonH)
	// JSON should do HTMLCharsAsIs by default
	testJsonH.HTMLCharsAsIs = true
	// set ExplicitRelease on each handle
	testMsgpackH.ExplicitRelease = true
	testBincH.ExplicitRelease = true
	testSimpleH.ExplicitRelease = true
	testCborH.ExplicitRelease = true
	testJsonH.ExplicitRelease = true

	testInitFlags()
	benchInitFlags()
}

func testInitFlags() {
	// delete(testDecOpts.ExtFuncs, timeTyp)
	flag.BoolVar(&testVerbose, "tv", false, "Text Extra Verbose Logging if -v if set")
	flag.BoolVar(&testInitDebug, "tg", false, "Test Init Debug")
	flag.IntVar(&testUseIoEncDec, "ti", -1, "Use IO Reader/Writer for Marshal/Unmarshal ie >= 0")
	flag.BoolVar(&testUseIoWrapper, "tiw", false, "Wrap the IO Reader/Writer with a base pass-through reader/writer")

	flag.BoolVar(&testSkipIntf, "tf", false, "Skip Interfaces")
	flag.BoolVar(&testUseReset, "tr", false, "Use Reset")
	flag.IntVar(&testNumRepeatString, "trs", 8, "Create string variables by repeating a string N times")
	flag.BoolVar(&testUseMust, "tm", true, "Use Must(En|De)code")

	flag.IntVar(&testMaxInitLen, "tx", 0, "Max Init Len")

	flag.IntVar(&testDepth, "tsd", 0, "Test Struc Depth")
	flag.BoolVar(&testMapStringKeyOnly, "tsk", false, "use maps with string keys only")
}

func benchInitFlags() {
	// flags reproduced here for compatibility (duplicate some in testInitFlags)
	flag.BoolVar(&testMapStringKeyOnly, "bs", false, "use maps with string keys only")
	flag.IntVar(&testDepth, "bd", 1, "Bench Depth")
}

func testHEDGet(h Handle) *testHED {
	for i := range testHEDs {
		v := &testHEDs[i]
		if v.H == h {
			return v
		}
	}
	testHEDs = append(testHEDs, testHED{h, NewEncoder(nil, h), NewDecoder(nil, h)})
	return &testHEDs[len(testHEDs)-1]
}

func testReinit() {
	testOnce = sync.Once{}
	testHEDs = nil
}

func testInitAll() {
	// only parse it once.
	if !flag.Parsed() {
		flag.Parse()
	}
	for _, f := range testPreInitFns {
		f()
	}
	for _, f := range testPostInitFns {
		f()
	}
}

func sTestCodecEncode(ts interface{}, bsIn []byte, fn func([]byte) *bytes.Buffer,
	h Handle, bh *BasicHandle) (bs []byte, err error) {
	// bs = make([]byte, 0, approxSize)
	var e *Encoder
	var buf *bytes.Buffer
	if testUseReset {
		e = testHEDGet(h).E
	} else {
		e = NewEncoder(nil, h)
	}
	var oldWriteBufferSize int
	if testUseIoEncDec >= 0 {
		buf = fn(bsIn)
		// set the encode options for using a buffer
		oldWriteBufferSize = bh.WriterBufferSize
		bh.WriterBufferSize = testUseIoEncDec
		if testUseIoWrapper {
			e.Reset(ioWriterWrapper{buf})
		} else {
			e.Reset(buf)
		}
	} else {
		bs = bsIn
		e.ResetBytes(&bs)
	}
	if testUseMust {
		e.MustEncode(ts)
	} else {
		err = e.Encode(ts)
	}
	if testUseIoEncDec >= 0 {
		bs = buf.Bytes()
		bh.WriterBufferSize = oldWriteBufferSize
	}
	if !testUseReset {
		e.Release()
	}
	return
}

func sTestCodecDecoder(bs []byte, h Handle, bh *BasicHandle) (d *Decoder, oldReadBufferSize int) {
	// var buf *bytes.Reader
	if testUseReset {
		d = testHEDGet(h).D
	} else {
		d = NewDecoder(nil, h)
	}
	if testUseIoEncDec >= 0 {
		buf := bytes.NewReader(bs)
		oldReadBufferSize = bh.ReaderBufferSize
		bh.ReaderBufferSize = testUseIoEncDec
		if testUseIoWrapper {
			d.Reset(ioReaderWrapper{buf})
		} else {
			d.Reset(buf)
		}
	} else {
		d.ResetBytes(bs)
	}
	return
}

func sTestCodecDecoderAfter(d *Decoder, oldReadBufferSize int, bh *BasicHandle) {
	if testUseIoEncDec >= 0 {
		bh.ReaderBufferSize = oldReadBufferSize
	}
	if !testUseReset {
		d.Release()
	}
}

func sTestCodecDecode(bs []byte, ts interface{}, h Handle, bh *BasicHandle) (err error) {
	d, oldReadBufferSize := sTestCodecDecoder(bs, h, bh)
	if testUseMust {
		d.MustDecode(ts)
	} else {
		err = d.Decode(ts)
	}
	sTestCodecDecoderAfter(d, oldReadBufferSize, bh)
	return
}

// // --- functions below are used by both benchmarks and tests

// // log message only when testVerbose = true (ie go test ... -- -tv).
// //
// // These are for intormational messages that do not necessarily
// // help with diagnosing a failure, or which are too large.
// func logTv(x interface{}, format string, args ...interface{}) {
// 	if !testVerbose {
// 		return
// 	}
// 	if t, ok := x.(testing.TB); ok { // only available from go 1.9
// 		t.Helper()
// 	}
// 	logT(x, format, args...)
// }

// // logT logs messages when running as go test -v
// //
// // Use it for diagnostics messages that help diagnost failure,
// // and when the output is not too long ie shorter than like 100 characters.
// //
// // In general, any logT followed by failT should call this.
// func logT(x interface{}, format string, args ...interface{}) {
// 	if x == nil {
// 		if len(format) == 0 || format[len(format)-1] != '\n' {
// 			format = format + "\n"
// 		}
// 		fmt.Printf(format, args...)
// 		return
// 	}
// 	if t, ok := x.(testing.TB); ok { // only available from go 1.9
// 		t.Helper()
// 		t.Logf(format, args...)
// 	}
// }

// func failTv(x testing.TB, args ...interface{}) {
// 	x.Helper()
// 	if testVerbose {
// 		failTMsg(x, args...)
// 	}
// 	x.FailNow()
// }

// func failT(x testing.TB, args ...interface{}) {
// 	x.Helper()
// 	failTMsg(x, args...)
// 	x.FailNow()
// }

// func failTMsg(x testing.TB, args ...interface{}) {
// 	x.Helper()
// 	if len(args) > 0 {
// 		if format, ok := args[0].(string); ok {
// 			logT(x, format, args[1:]...)
// 		} else if len(args) == 1 {
// 			logT(x, "%v", args[0])
// 		} else {
// 			logT(x, "%v", args)
// 		}
// 	}
// }

// --- functions below are used only by benchmarks alone

func fnBenchmarkByteBuf(bsIn []byte) (buf *bytes.Buffer) {
	// var buf bytes.Buffer
	// buf.Grow(approxSize)
	buf = bytes.NewBuffer(bsIn)
	buf.Truncate(0)
	return
}

// func benchFnCodecEncode(ts interface{}, bsIn []byte, h Handle) (bs []byte, err error) {
// 	return testCodecEncode(ts, bsIn, fnBenchmarkByteBuf, h)
// }

// func benchFnCodecDecode(bs []byte, ts interface{}, h Handle) (err error) {
// 	return testCodecDecode(bs, ts, h)
// }
<commit_msg>update shared_test.go<commit_after>// Copyright (c) 2012-2020 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a MIT license found in the LICENSE file.

package codec

// This file sets up the variables used, including testInitFns.
// Each file should add initialization that should be performed
// after flags are parsed.
//
// init is a multi-step process:
//   - setup vars (handled by init functions in each file)
//   - parse flags
//   - setup derived vars (handled by pre-init registered functions - registered in init function)
//   - post init (handled by post-init registered functions - registered in init function)
// This way, no one has to manage carefully control the initialization
// using file names, etc.
//
// Tests which require external dependencies need the -tag=x parameter.
// They should be run as:
//    go test -tags=x -run=. <other parameters ...>
// Benchmarks should also take this parameter, to include the sereal, xdr, etc.
// To run against codecgen, etc, make sure you pass extra parameters.
// Example usage:
//    go test "-tags=x codecgen" -bench=. <other parameters ...>
//
// To fully test everything:
//    go test -tags=x -benchtime=100ms -tv -bg -bi  -brw -bu -v -run=. -bench=.

// Handling flags
// codec_test.go will define a set of global flags for testing, including:
//   - Use Reset
//   - Use IO reader/writer (vs direct bytes)
//   - Set Canonical
//   - Set InternStrings
//   - Use Symbols
//
// This way, we can test them all by running same set of tests with a different
// set of flags.
//
// Following this, all the benchmarks will utilize flags set by codec_test.go
// and will not redefine these "global" flags.

import (
	"bytes"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"sync"
)

import . "github.com/ugorji/go/codec"

type testHED struct {
	H Handle
	E *Encoder
	D *Decoder
}

type ioReaderWrapper struct {
	r io.Reader
}

func (x ioReaderWrapper) Read(p []byte) (n int, err error) {
	return x.r.Read(p)
}

type ioWriterWrapper struct {
	w io.Writer
}

func (x ioWriterWrapper) Write(p []byte) (n int, err error) {
	return x.w.Write(p)
}

var (
	// testNoopH    = NoopHandle(8)
	testMsgpackH = &MsgpackHandle{}
	testBincH    = &BincHandle{}
	testSimpleH  = &SimpleHandle{}
	testCborH    = &CborHandle{}
	testJsonH    = &JsonHandle{}

	testHandles     []Handle
	testPreInitFns  []func()
	testPostInitFns []func()

	testOnce sync.Once

	testHEDs []testHED
)

// flag variables used by tests (and bench)
var (
	testVerbose bool

	//depth of 0 maps to ~400bytes json-encoded string, 1 maps to ~1400 bytes, etc
	//For depth>1, we likely trigger stack growth for encoders, making benchmarking unreliable.
	testDepth int

	testMaxInitLen int

	testInitDebug bool
	testUseReset  bool
	testSkipIntf  bool
	testUseMust   bool

	testUseIoEncDec  int
	testUseIoWrapper bool

	testNumRepeatString int

	testRpcBufsize       int
	testMapStringKeyOnly bool
)

// variables that are not flags, but which can configure the handles
var (
	testEncodeOptions EncodeOptions
	testDecodeOptions DecodeOptions
	testRPCOptions    RPCOptions
)

func init() {
	log.SetOutput(ioutil.Discard) // don't allow things log to standard out/err
	testHEDs = make([]testHED, 0, 32)
	testHandles = append(testHandles,
		// testNoopH,
		testMsgpackH, testBincH, testSimpleH, testCborH, testJsonH)
	// JSON should do HTMLCharsAsIs by default
	testJsonH.HTMLCharsAsIs = true
	// set ExplicitRelease on each handle
	testMsgpackH.ExplicitRelease = true
	testBincH.ExplicitRelease = true
	testSimpleH.ExplicitRelease = true
	testCborH.ExplicitRelease = true
	testJsonH.ExplicitRelease = true

	testInitFlags()
	benchInitFlags()
}

func testInitFlags() {
	// delete(testDecOpts.ExtFuncs, timeTyp)
	flag.BoolVar(&testVerbose, "tv", false, "Text Extra Verbose Logging if -v if set")
	flag.BoolVar(&testInitDebug, "tg", false, "Test Init Debug")
	flag.IntVar(&testUseIoEncDec, "ti", -1, "Use IO Reader/Writer for Marshal/Unmarshal ie >= 0")
	flag.BoolVar(&testUseIoWrapper, "tiw", false, "Wrap the IO Reader/Writer with a base pass-through reader/writer")

	flag.BoolVar(&testSkipIntf, "tf", false, "Skip Interfaces")
	flag.BoolVar(&testUseReset, "tr", false, "Use Reset")
	flag.IntVar(&testNumRepeatString, "trs", 8, "Create string variables by repeating a string N times")
	flag.BoolVar(&testUseMust, "tm", true, "Use Must(En|De)code")

	flag.IntVar(&testMaxInitLen, "tx", 0, "Max Init Len")

	flag.IntVar(&testDepth, "tsd", 0, "Test Struc Depth")
	flag.BoolVar(&testMapStringKeyOnly, "tsk", false, "use maps with string keys only")
}

func benchInitFlags() {
	// flags reproduced here for compatibility (duplicate some in testInitFlags)
	flag.BoolVar(&testMapStringKeyOnly, "bs", false, "use maps with string keys only")
	flag.IntVar(&testDepth, "bd", 1, "Bench Depth")
}

func testHEDGet(h Handle) *testHED {
	for i := range testHEDs {
		v := &testHEDs[i]
		if v.H == h {
			return v
		}
	}
	testHEDs = append(testHEDs, testHED{h, NewEncoder(nil, h), NewDecoder(nil, h)})
	return &testHEDs[len(testHEDs)-1]
}

func testReinit() {
	testOnce = sync.Once{}
	testHEDs = nil
}

func testInitAll() {
	// only parse it once.
	if !flag.Parsed() {
		flag.Parse()
	}
	for _, f := range testPreInitFns {
		f()
	}
	for _, f := range testPostInitFns {
		f()
	}
}

func sTestCodecEncode(ts interface{}, bsIn []byte, fn func([]byte) *bytes.Buffer,
	h Handle, bh *BasicHandle) (bs []byte, err error) {
	// bs = make([]byte, 0, approxSize)
	var e *Encoder
	var buf *bytes.Buffer
	if testUseReset {
		e = testHEDGet(h).E
	} else {
		e = NewEncoder(nil, h)
	}
	var oldWriteBufferSize int
	if testUseIoEncDec >= 0 {
		buf = fn(bsIn)
		// set the encode options for using a buffer
		oldWriteBufferSize = bh.WriterBufferSize
		bh.WriterBufferSize = testUseIoEncDec
		if testUseIoWrapper {
			e.Reset(ioWriterWrapper{buf})
		} else {
			e.Reset(buf)
		}
	} else {
		bs = bsIn
		e.ResetBytes(&bs)
	}
	if testUseMust {
		e.MustEncode(ts)
	} else {
		err = e.Encode(ts)
	}
	if testUseIoEncDec >= 0 {
		bs = buf.Bytes()
		bh.WriterBufferSize = oldWriteBufferSize
	}
	if !testUseReset {
		e.Release()
	}
	return
}

func sTestCodecDecoder(bs []byte, h Handle, bh *BasicHandle) (d *Decoder, oldReadBufferSize int) {
	// var buf *bytes.Reader
	if testUseReset {
		d = testHEDGet(h).D
	} else {
		d = NewDecoder(nil, h)
	}
	if testUseIoEncDec >= 0 {
		buf := bytes.NewReader(bs)
		oldReadBufferSize = bh.ReaderBufferSize
		bh.ReaderBufferSize = testUseIoEncDec
		if testUseIoWrapper {
			d.Reset(ioReaderWrapper{buf})
		} else {
			d.Reset(buf)
		}
	} else {
		d.ResetBytes(bs)
	}
	return
}

func sTestCodecDecoderAfter(d *Decoder, oldReadBufferSize int, bh *BasicHandle) {
	if testUseIoEncDec >= 0 {
		bh.ReaderBufferSize = oldReadBufferSize
	}
	if !testUseReset {
		d.Release()
	}
}

func sTestCodecDecode(bs []byte, ts interface{}, h Handle, bh *BasicHandle) (err error) {
	d, oldReadBufferSize := sTestCodecDecoder(bs, h, bh)
	if testUseMust {
		d.MustDecode(ts)
	} else {
		err = d.Decode(ts)
	}
	sTestCodecDecoderAfter(d, oldReadBufferSize, bh)
	return
}

// // --- functions below are used by both benchmarks and tests

// // log message only when testVerbose = true (ie go test ... -- -tv).
// //
// // These are for intormational messages that do not necessarily
// // help with diagnosing a failure, or which are too large.
// func logTv(x interface{}, format string, args ...interface{}) {
// 	if !testVerbose {
// 		return
// 	}
// 	if t, ok := x.(testing.TB); ok { // only available from go 1.9
// 		t.Helper()
// 	}
// 	logT(x, format, args...)
// }

// // logT logs messages when running as go test -v
// //
// // Use it for diagnostics messages that help diagnost failure,
// // and when the output is not too long ie shorter than like 100 characters.
// //
// // In general, any logT followed by failT should call this.
// func logT(x interface{}, format string, args ...interface{}) {
// 	if x == nil {
// 		if len(format) == 0 || format[len(format)-1] != '\n' {
// 			format = format + "\n"
// 		}
// 		fmt.Printf(format, args...)
// 		return
// 	}
// 	if t, ok := x.(testing.TB); ok { // only available from go 1.9
// 		t.Helper()
// 		t.Logf(format, args...)
// 	}
// }

// func failTv(x testing.TB, args ...interface{}) {
// 	x.Helper()
// 	if testVerbose {
// 		failTMsg(x, args...)
// 	}
// 	x.FailNow()
// }

// func failT(x testing.TB, args ...interface{}) {
// 	x.Helper()
// 	failTMsg(x, args...)
// 	x.FailNow()
// }

// func failTMsg(x testing.TB, args ...interface{}) {
// 	x.Helper()
// 	if len(args) > 0 {
// 		if format, ok := args[0].(string); ok {
// 			logT(x, format, args[1:]...)
// 		} else if len(args) == 1 {
// 			logT(x, "%v", args[0])
// 		} else {
// 			logT(x, "%v", args)
// 		}
// 	}
// }

// --- functions below are used only by benchmarks alone

func fnBenchmarkByteBuf(bsIn []byte) (buf *bytes.Buffer) {
	// var buf bytes.Buffer
	// buf.Grow(approxSize)
	buf = bytes.NewBuffer(bsIn)
	buf.Truncate(0)
	return
}

// func benchFnCodecEncode(ts interface{}, bsIn []byte, h Handle) (bs []byte, err error) {
// 	return testCodecEncode(ts, bsIn, fnBenchmarkByteBuf, h)
// }

// func benchFnCodecDecode(bs []byte, ts interface{}, h Handle) (err error) {
// 	return testCodecDecode(bs, ts, h)
// }
<|endoftext|>


<commit_before>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// thriftrw-plugin-yarpc implements a plugin for ThriftRW that generates code
// compatible with YARPC.
//
// For more information, check the documentation of the parent package.
package main

import (
	"path/filepath"
	"strings"

	"go.uber.org/thriftrw/plugin"
	"go.uber.org/thriftrw/plugin/api"
)

const serverTemplate = `
// Code generated by thriftrw-plugin-yarpc
// @generated

<$pkgname := printf "%sserver" (lower .Service.Name)>
package <$pkgname>

<$yarpc     := import "go.uber.org/yarpc">
<$thrift    := import "go.uber.org/yarpc/encoding/thrift">
<$transport := import "go.uber.org/yarpc/api/transport">
<$context   := import "context">

// Interface is the server-side interface for the <.Service.Name> service.
type Interface interface {
	<if .Parent>
		<$parentPath := printf "%s/yarpc/%sserver" .ParentModule.ImportPath (lower .Parent.Name)>
		<import $parentPath>.Interface
	<end>

	<range .Service.Functions>
		<.Name>(
			ctx <$context>.Context,
			reqMeta <$yarpc>.ReqMeta, <range .Arguments>
			<.Name> <formatType .Type>,<end>
		)<if .OneWay> error
		<else if .ReturnType> (<formatType .ReturnType>, <$yarpc>.ResMeta, error)
		<else> (<$yarpc>.ResMeta, error)
		<end>
	<end>
}

// New prepares an implementation of the <.Service.Name> service for
// registration.
//
// 	handler := <.Service.Name>Handler{}
// 	dispatcher.Register(<$pkgname>.New(handler))
func New(impl Interface, opts ...<$thrift>.RegisterOption) []<$transport>.Registrant {
	h := handler{impl}
	service := <$thrift>.Service{
		Name: "<.Service.Name>",
		Methods: map[string]<$thrift>.UnaryHandler{
			<range .Service.Functions><if not .OneWay>"<.ThriftName>": <$thrift>.UnaryHandlerFunc(h.<.Name>),<end>
		<end>},
		OnewayMethods: map[string]<$thrift>.OnewayHandler{
			<range .Service.Functions><if .OneWay>"<.ThriftName>": <$thrift>.OnewayHandlerFunc(h.<.Name>),<end>
		<end>},
	}
	return <$thrift>.BuildRegistrants(service, opts...)
}

type handler struct{ impl Interface }

<$service := .Service>
<$module := .Module>
<range .Service.Functions>
<$prefix := printf "%s.%s_%s_" (import $module.ImportPath) $service.Name .Name>

<$wire := import "go.uber.org/thriftrw/wire">

<if .OneWay>
func (h handler) <.Name>(
	ctx <$context>.Context,
	reqMeta <$yarpc>.ReqMeta,
	body <$wire>.Value,
) error {
	var args <$prefix>Args
	if err := args.FromWire(body); err != nil {
		return err
	}

	return h.impl.<.Name>(ctx, reqMeta, <range .Arguments>args.<.Name>,<end>)
}
<else>
func (h handler) <.Name>(
	ctx <$context>.Context,
	reqMeta <$yarpc>.ReqMeta,
	body <$wire>.Value,
) (<$thrift>.Response, error) {
	var args <$prefix>Args
	if err := args.FromWire(body); err != nil {
		return <$thrift>.Response{}, err
	}

	<if .ReturnType>
		success, resMeta, err := h.impl.<.Name>(ctx, reqMeta, <range .Arguments>args.<.Name>,<end>)
	<else>
		resMeta, err := h.impl.<.Name>(ctx, reqMeta, <range .Arguments>args.<.Name>,<end>)
	<end>

	hadError := err != nil
	result, err := <$prefix>Helper.WrapResponse(<if .ReturnType>success,<end> err)

	var response <$thrift>.Response
	if err == nil {
		response.IsApplicationError = hadError
		response.Meta = resMeta
		response.Body = result
	}
	return response, err
}
<end>
<end>
`

const clientTemplate = `
// Code generated by thriftrw-plugin-yarpc
// @generated

<$pkgname := printf "%sclient" (lower .Service.Name)>
package <$pkgname>

<$yarpc     := import "go.uber.org/yarpc">
<$transport := import "go.uber.org/yarpc/api/transport">
<$thrift    := import "go.uber.org/yarpc/encoding/thrift">
<$context   := import "context">

// Interface is a client for the <.Service.Name> service.
type Interface interface {
	<if .Parent>
		<$parentPath := printf "%s/yarpc/%sclient" .ParentModule.ImportPath (lower .Parent.Name)>
		<import $parentPath>.Interface
	<end>

	<range .Service.Functions>
		<.Name>(
			ctx <$context>.Context,
			reqMeta <$yarpc>.CallReqMeta, <range .Arguments>
				<.Name> <formatType .Type>,<end>
		)<if .OneWay> (<$yarpc>.Ack, error)
		<else if .ReturnType> (<formatType .ReturnType>, <$yarpc>.CallResMeta, error)
		<else> (<$yarpc>.CallResMeta, error)
		<end>
	<end>
}

</* TODO(abg): Pull the default routing name from a Thrift annotation? */>

// New builds a new client for the <.Service.Name> service.
//
// 	client := <$pkgname>.New(dispatcher.ClientConfig("<lower .Service.Name>"))
func New(c <$transport>.ClientConfig, opts ...<$thrift>.ClientOption) Interface {
	return client{c: <$thrift>.New(<$thrift>.Config{
		Service: "<.Service.Name>",
		ClientConfig: c,
	}, opts...)}
}

func init() {
	<$yarpc>.RegisterClientBuilder(func(c <$transport>.ClientConfig) Interface {
		return New(c)
	})
}

type client struct{ c <$thrift>.Client }

<$service := .Service>
<$module := .Module>
<range .Service.Functions>
<$prefix := printf "%s.%s_%s_" (import $module.ImportPath) $service.Name .Name>

func (c client) <.Name>(
	ctx <$context>.Context,
	reqMeta <$yarpc>.CallReqMeta, <range .Arguments>
	_<.Name> <formatType .Type>,<end>
<if .OneWay>) (<$yarpc>.Ack, error) {
	args := <$prefix>Helper.Args(<range .Arguments>_<.Name>, <end>)
	return c.c.CallOneway(ctx, reqMeta, args)
}
<else>) (<if .ReturnType>success <formatType .ReturnType>,<end> resMeta <$yarpc>.CallResMeta, err error) {
	<$wire := import "go.uber.org/thriftrw/wire">
	args := <$prefix>Helper.Args(<range .Arguments>_<.Name>, <end>)

	var body <$wire>.Value
	body, resMeta, err = c.c.Call(ctx, reqMeta, args)
	if err != nil {
		return
	}

	var result <$prefix>Result
	if err = result.FromWire(body); err != nil {
		return
	}

	<if .ReturnType>success, <end>err = <$prefix>Helper.UnwrapResponse(&result)
	return
}
<end>
<end>
`

var templateOptions = []plugin.TemplateOption{
	plugin.TemplateFunc("lower", strings.ToLower),
}

type generator struct{}

func (generator) Generate(req *api.GenerateServiceRequest) (*api.GenerateServiceResponse, error) {
	files := make(map[string][]byte)
	for _, serviceID := range req.RootServices {
		service := req.Services[serviceID]
		module := req.Modules[service.ModuleID]

		var (
			parent       *api.Service
			parentModule *api.Module
		)
		if service.ParentID != nil {
			parent = req.Services[*service.ParentID]
			parentModule = req.Modules[parent.ModuleID]
		}

		templateData := struct {
			Module       *api.Module
			Service      *api.Service
			Parent       *api.Service
			ParentModule *api.Module
		}{
			Module:       module,
			Service:      service,
			Parent:       parent,
			ParentModule: parentModule,
		}

		// kv.thrift => .../kv/yarpc
		baseDir := filepath.Join(module.Directory, "yarpc")

		serverPackageName := strings.ToLower(service.Name) + "server"
		clientPackageName := strings.ToLower(service.Name) + "client"

		// kv.thrift =>
		//   .../yarpc/keyvalueserver/server.go
		//   .../yarpc/keyvalueclient/client.go
		serverFilePath := filepath.Join(baseDir, serverPackageName, "server.go")
		clientFilePath := filepath.Join(baseDir, clientPackageName, "client.go")

		serverContents, err := plugin.GoFileFromTemplate(
			serverFilePath, serverTemplate, templateData, templateOptions...)
		if err != nil {
			return nil, err
		}

		clientContents, err := plugin.GoFileFromTemplate(
			clientFilePath, clientTemplate, templateData, templateOptions...)
		if err != nil {
			return nil, err
		}

		files[serverFilePath] = serverContents
		files[clientFilePath] = clientContents
	}
	return &api.GenerateServiceResponse{Files: files}, nil
}

func main() {
	plugin.Main(&plugin.Plugin{Name: "yarpc", ServiceGenerator: generator{}})
}
<commit_msg>Parameterize ThriftRW plugin (#582)<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

// thriftrw-plugin-yarpc implements a plugin for ThriftRW that generates code
// compatible with YARPC.
//
// For more information, check the documentation of the parent package.
package main

import (
	"flag"
	"path/filepath"
	"strings"

	"go.uber.org/thriftrw/plugin"
	"go.uber.org/thriftrw/plugin/api"
)

var (
	_context = flag.String("context-import-path",
		"context",
		"Import path at which Context is available")
	_unaryHandlerWrapper = flag.String("unary-handler-wrapper",
		"go.uber.org/yarpc/encoding/thrift.UnaryHandlerFunc",
		"Function used to wrap generic Thrift unary function handlers into YARPC handlers")
	_onewayHandlerWrapper = flag.String("oneway-handler-wrapper",
		"go.uber.org/yarpc/encoding/thrift.OnewayHandlerFunc",
		"Function used to wrap generic Thrift oneway function handlers into YARPC handlers")
)

const serverTemplate = `
// Code generated by thriftrw-plugin-yarpc
// @generated

<$pkgname := printf "%sserver" (lower .Service.Name)>
package <$pkgname>
<$yarpc     := import "go.uber.org/yarpc">
<$thrift    := import "go.uber.org/yarpc/encoding/thrift">
<$transport := import "go.uber.org/yarpc/api/transport">
<$context   := import .ContextImportPath>

// Interface is the server-side interface for the <.Service.Name> service.
type Interface interface {
	<if .Parent>
		<$parentPath := printf "%s/yarpc/%sserver" .ParentModule.ImportPath (lower .Parent.Name)>
		<import $parentPath>.Interface
	<end>

	<range .Service.Functions>
		<.Name>(
			ctx <$context>.Context,
			reqMeta <$yarpc>.ReqMeta, <range .Arguments>
			<.Name> <formatType .Type>,<end>
		)<if .OneWay> error
		<else if .ReturnType> (<formatType .ReturnType>, <$yarpc>.ResMeta, error)
		<else> (<$yarpc>.ResMeta, error)
		<end>
	<end>
}

// New prepares an implementation of the <.Service.Name> service for
// registration.
//
// 	handler := <.Service.Name>Handler{}
// 	dispatcher.Register(<$pkgname>.New(handler))
func New(impl Interface, opts ...<$thrift>.RegisterOption) []<$transport>.Registrant {
	h := handler{impl}
	service := <$thrift>.Service{
		Name: "<.Service.Name>",
			Methods: map[string]<$thrift>.UnaryHandler{
				<$unaryWrapperImport := .UnaryWrapperImport>
				<$unaryWrapperFunc := .UnaryWrapperFunc>
				<range .Service.Functions>
					<if not .OneWay>"<.ThriftName>": <import $unaryWrapperImport>.<$unaryWrapperFunc>(h.<.Name>),<end>
			<end>},
			OnewayMethods: map[string]<$thrift>.OnewayHandler{
				<$onewayWrapperImport := .OnewayWrapperImport>
				<$onewayWrapperFunc := .OnewayWrapperFunc>
				<range .Service.Functions>
					<if .OneWay>"<.ThriftName>": <import $onewayWrapperImport>.<$onewayWrapperFunc>(h.<.Name>),<end>
			<end>},
	}
	return <$thrift>.BuildRegistrants(service, opts...)
}

type handler struct{ impl Interface }

<$service := .Service>
<$module := .Module>
<range .Service.Functions>
<$prefix := printf "%s.%s_%s_" (import $module.ImportPath) $service.Name .Name>

<$wire := import "go.uber.org/thriftrw/wire">

<if .OneWay>
func (h handler) <.Name>(
	ctx <$context>.Context,
	reqMeta <$yarpc>.ReqMeta,
	body <$wire>.Value,
) error {
	var args <$prefix>Args
	if err := args.FromWire(body); err != nil {
		return err
	}

	return h.impl.<.Name>(ctx, reqMeta, <range .Arguments>args.<.Name>,<end>)
}
<else>
func (h handler) <.Name>(
	ctx <$context>.Context,
	reqMeta <$yarpc>.ReqMeta,
	body <$wire>.Value,
) (<$thrift>.Response, error) {
	var args <$prefix>Args
	if err := args.FromWire(body); err != nil {
		return <$thrift>.Response{}, err
	}

	<if .ReturnType>
		success, resMeta, err := h.impl.<.Name>(ctx, reqMeta, <range .Arguments>args.<.Name>,<end>)
	<else>
		resMeta, err := h.impl.<.Name>(ctx, reqMeta, <range .Arguments>args.<.Name>,<end>)
	<end>

	hadError := err != nil
	result, err := <$prefix>Helper.WrapResponse(<if .ReturnType>success,<end> err)

	var response <$thrift>.Response
	if err == nil {
		response.IsApplicationError = hadError
		response.Meta = resMeta
		response.Body = result
	}
	return response, err
}
<end>
<end>
`

const clientTemplate = `
// Code generated by thriftrw-plugin-yarpc
// @generated

<$pkgname := printf "%sclient" (lower .Service.Name)>
package <$pkgname>

<$yarpc     := import "go.uber.org/yarpc">
<$transport := import "go.uber.org/yarpc/api/transport">
<$thrift    := import "go.uber.org/yarpc/encoding/thrift">
<$context   := import "context">

// Interface is a client for the <.Service.Name> service.
type Interface interface {
	<if .Parent>
		<$parentPath := printf "%s/yarpc/%sclient" .ParentModule.ImportPath (lower .Parent.Name)>
		<import $parentPath>.Interface
	<end>

	<range .Service.Functions>
		<.Name>(
			ctx <$context>.Context,
			reqMeta <$yarpc>.CallReqMeta, <range .Arguments>
				<.Name> <formatType .Type>,<end>
		)<if .OneWay> (<$yarpc>.Ack, error)
		<else if .ReturnType> (<formatType .ReturnType>, <$yarpc>.CallResMeta, error)
		<else> (<$yarpc>.CallResMeta, error)
		<end>
	<end>
}

</* TODO(abg): Pull the default routing name from a Thrift annotation? */>

// New builds a new client for the <.Service.Name> service.
//
// 	client := <$pkgname>.New(dispatcher.ClientConfig("<lower .Service.Name>"))
func New(c <$transport>.ClientConfig, opts ...<$thrift>.ClientOption) Interface {
	return client{c: <$thrift>.New(<$thrift>.Config{
		Service: "<.Service.Name>",
		ClientConfig: c,
	}, opts...)}
}

func init() {
	<$yarpc>.RegisterClientBuilder(func(c <$transport>.ClientConfig) Interface {
		return New(c)
	})
}

type client struct{ c <$thrift>.Client }

<$service := .Service>
<$module := .Module>
<range .Service.Functions>
<$prefix := printf "%s.%s_%s_" (import $module.ImportPath) $service.Name .Name>

func (c client) <.Name>(
	ctx <$context>.Context,
	reqMeta <$yarpc>.CallReqMeta, <range .Arguments>
	_<.Name> <formatType .Type>,<end>
<if .OneWay>) (<$yarpc>.Ack, error) {
	args := <$prefix>Helper.Args(<range .Arguments>_<.Name>, <end>)
	return c.c.CallOneway(ctx, reqMeta, args)
}
<else>) (<if .ReturnType>success <formatType .ReturnType>,<end> resMeta <$yarpc>.CallResMeta, err error) {
	<$wire := import "go.uber.org/thriftrw/wire">
	args := <$prefix>Helper.Args(<range .Arguments>_<.Name>, <end>)

	var body <$wire>.Value
	body, resMeta, err = c.c.Call(ctx, reqMeta, args)
	if err != nil {
		return
	}

	var result <$prefix>Result
	if err = result.FromWire(body); err != nil {
		return
	}

	<if .ReturnType>success, <end>err = <$prefix>Helper.UnwrapResponse(&result)
	return
}
<end>
<end>
`

var templateOptions = []plugin.TemplateOption{
	plugin.TemplateFunc("lower", strings.ToLower),
}

type generator struct{}

func (generator) Generate(req *api.GenerateServiceRequest) (*api.GenerateServiceResponse, error) {
	files := make(map[string][]byte)
	for _, serviceID := range req.RootServices {
		service := req.Services[serviceID]
		module := req.Modules[service.ModuleID]

		var (
			parent       *api.Service
			parentModule *api.Module
		)
		if service.ParentID != nil {
			parent = req.Services[*service.ParentID]
			parentModule = req.Modules[parent.ModuleID]
		}

		unaryWrapperImport, unaryWrapperFunc := splitFunctionPath(*_unaryHandlerWrapper)
		onewayWrapperImport, onewayWrapperFunc := splitFunctionPath(*_onewayHandlerWrapper)

		templateData := struct {
			Module       *api.Module
			Service      *api.Service
			Parent       *api.Service
			ParentModule *api.Module

			ContextImportPath   string
			UnaryWrapperImport  string
			UnaryWrapperFunc    string
			OnewayWrapperImport string
			OnewayWrapperFunc   string
		}{
			Module:       module,
			Service:      service,
			Parent:       parent,
			ParentModule: parentModule,

			ContextImportPath:   *_context,
			UnaryWrapperImport:  unaryWrapperImport,
			UnaryWrapperFunc:    unaryWrapperFunc,
			OnewayWrapperImport: onewayWrapperImport,
			OnewayWrapperFunc:   onewayWrapperFunc,
		}

		// kv.thrift => .../kv/yarpc
		baseDir := filepath.Join(module.Directory, "yarpc")

		serverPackageName := strings.ToLower(service.Name) + "server"
		clientPackageName := strings.ToLower(service.Name) + "client"

		// kv.thrift =>
		//   .../yarpc/keyvalueserver/server.go
		//   .../yarpc/keyvalueclient/client.go
		serverFilePath := filepath.Join(baseDir, serverPackageName, "server.go")
		clientFilePath := filepath.Join(baseDir, clientPackageName, "client.go")

		serverContents, err := plugin.GoFileFromTemplate(
			serverFilePath, serverTemplate, templateData, templateOptions...)
		if err != nil {
			return nil, err
		}

		clientContents, err := plugin.GoFileFromTemplate(
			clientFilePath, clientTemplate, templateData, templateOptions...)
		if err != nil {
			return nil, err
		}

		files[serverFilePath] = serverContents
		files[clientFilePath] = clientContents
	}
	return &api.GenerateServiceResponse{Files: files}, nil
}

func splitFunctionPath(input string) (string, string) {
	i := strings.LastIndex(input, ".")
	return input[:i], input[i+1:]
}

func main() {
	flag.Parse()
	plugin.Main(&plugin.Plugin{Name: "yarpc", ServiceGenerator: generator{}})
}
<|endoftext|>