<commit_before>package main

import (
	"archive/tar"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/moby/src/initrd"
)

// Process the build arguments and execute build
func build(args []string) {
	buildCmd := flag.NewFlagSet("build", flag.ExitOnError)
	buildCmd.Usage = func() {
		fmt.Printf("USAGE: %s build [options] [file.yml]\n\n", os.Args[0])
		fmt.Printf("'file.yml' defaults to 'moby.yml' if not specified.\n\n")
		fmt.Printf("Options:\n")
		buildCmd.PrintDefaults()
	}
	buildName := buildCmd.String("name", "", "Name to use for output files")
	buildPull := buildCmd.Bool("pull", false, "Always pull images")

	buildCmd.Parse(args)
	remArgs := buildCmd.Args()

	conf := "moby.yml"
	if len(remArgs) > 0 {
		conf = remArgs[0]
	}

	buildInternal(*buildName, *buildPull, conf)
}

// Perform the actual build process
func buildInternal(name string, pull bool, conf string) {
	if name == "" {
		name = filepath.Base(conf)
		ext := filepath.Ext(conf)
		if ext != "" {
			name = name[:len(name)-len(ext)]
		}
	}

	config, err := ioutil.ReadFile(conf)
	if err != nil {
		log.Fatalf("Cannot open config file: %v", err)
	}

	m, err := NewConfig(config)
	if err != nil {
		log.Fatalf("Invalid config: %v", err)
	}

	containers := []*bytes.Buffer{}

	if pull {
		log.Infof("Pull kernel image: %s", m.Kernel.Image)
		err := dockerPull(m.Kernel.Image)
		if err != nil {
			log.Fatalf("Could not pull image %s: %v", m.Kernel.Image, err)
		}
	}
	// get kernel bzImage and initrd tarball from container
	// TODO examine contents to see what names they might have
	log.Infof("Extract kernel image: %s", m.Kernel.Image)
	const (
		bzimageName = "bzImage"
		ktarName    = "kernel.tar"
	)
	out, err := dockerRun(m.Kernel.Image, "tar", "cf", "-", bzimageName, ktarName)
	if err != nil {
		log.Fatalf("Failed to extract kernel image and tarball: %v", err)
	}
	buf := bytes.NewBuffer(out)
	bzimage, ktar, err := untarKernel(buf, bzimageName, ktarName)
	if err != nil {
		log.Fatalf("Could not extract bzImage and kernel filesystem from tarball. %v", err)
	}
	containers = append(containers, ktar)

	// convert init image to tarball
	if pull {
		log.Infof("Pull init: %s", m.Init)
		err := dockerPull(m.Init)
		if err != nil {
			log.Fatalf("Could not pull image %s: %v", m.Init, err)
		}
	}
	log.Infof("Process init: %s", m.Init)
	init, err := ImageExtract(m.Init, "")
	if err != nil {
		log.Fatalf("Failed to build init tarball: %v", err)
	}
	buffer := bytes.NewBuffer(init)
	containers = append(containers, buffer)

	log.Infof("Add system containers:")
	for i, image := range m.System {
		if pull {
			log.Infof("  Pull: %s", image.Image)
			err := dockerPull(image.Image)
			if err != nil {
				log.Fatalf("Could not pull image %s: %v", image.Image, err)
			}
		}
		log.Infof("  Create OCI config for %s", image.Image)
		config, err := ConfigToOCI(&image)
		if err != nil {
			log.Fatalf("Failed to run riddler to get config.json for %s: %v", image.Image, err)
		}
		so := fmt.Sprintf("%03d", i)
		path := "containers/system/" + so + "-" + image.Name
		out, err := ImageBundle(path, image.Image, config)
		if err != nil {
			log.Fatalf("Failed to extract root filesystem for %s: %v", image.Image, err)
		}
		buffer := bytes.NewBuffer(out)
		containers = append(containers, buffer)
	}

	log.Infof("Add daemon containers:")
	for _, image := range m.Daemon {
		if pull {
			log.Infof("  Pull: %s", image.Image)
			err := dockerPull(image.Image)
			if err != nil {
				log.Fatalf("Could not pull image %s: %v", image.Image, err)
			}
		}
		log.Infof("  Create OCI config for %s", image.Image)
		config, err := ConfigToOCI(&image)
		if err != nil {
			log.Fatalf("Failed to run riddler to get config.json for %s: %v", image.Image, err)
		}
		path := "containers/daemon/" + image.Name
		out, err := ImageBundle(path, image.Image, config)
		if err != nil {
			log.Fatalf("Failed to extract root filesystem for %s: %v", image.Image, err)
		}
		buffer := bytes.NewBuffer(out)
		containers = append(containers, buffer)
	}

	// add files
	buffer, err = filesystem(m)
	if err != nil {
		log.Fatalf("failed to add filesystem parts: %v", err)
	}
	containers = append(containers, buffer)

	log.Infof("Create initial ram disk")
	initrd, err := containersInitrd(containers)
	if err != nil {
		log.Fatalf("Failed to make initrd %v", err)
	}

	log.Infof("Create outputs:")
	err = outputs(m, name, bzimage.Bytes(), initrd.Bytes())
	if err != nil {
		log.Fatalf("Error writing outputs: %v", err)
	}
}

func untarKernel(buf *bytes.Buffer, bzimageName, ktarName string) (*bytes.Buffer, *bytes.Buffer, error) {
	tr := tar.NewReader(buf)

	var bzimage, ktar *bytes.Buffer

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalln(err)
		}
		switch hdr.Name {
		case bzimageName:
			bzimage = new(bytes.Buffer)
			_, err := io.Copy(bzimage, tr)
			if err != nil {
				return nil, nil, err
			}
		case ktarName:
			ktar = new(bytes.Buffer)
			_, err := io.Copy(bzimage, tr)
			if err != nil {
				return nil, nil, err
			}
		default:
			continue
		}
	}

	if ktar == nil || bzimage == nil {
		return nil, nil, errors.New("did not find bzImage and kernel.tar in tarball")
	}

	return bzimage, ktar, nil
}

func containersInitrd(containers []*bytes.Buffer) (*bytes.Buffer, error) {
	w := new(bytes.Buffer)
	iw := initrd.NewWriter(w)
	defer iw.Close()
	for _, file := range containers {
		_, err := initrd.Copy(iw, file)
		if err != nil {
			return nil, err
		}
	}

	return w, nil
}
<commit_msg>cli: Don't default to moby.yml<commit_after>package main

import (
	"archive/tar"
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/moby/src/initrd"
)

// Process the build arguments and execute build
func build(args []string) {
	buildCmd := flag.NewFlagSet("build", flag.ExitOnError)
	buildCmd.Usage = func() {
		fmt.Printf("USAGE: %s build [options] <file.yml>\n\n", os.Args[0])
		fmt.Printf("Options:\n")
		buildCmd.PrintDefaults()
	}
	buildName := buildCmd.String("name", "", "Name to use for output files")
	buildPull := buildCmd.Bool("pull", false, "Always pull images")

	buildCmd.Parse(args)
	remArgs := buildCmd.Args()

	if len(remArgs) == 0 {
		fmt.Println("Please specify a configuration file\n")
		buildCmd.Usage()
		os.Exit(1)
	}
	conf := remArgs[0]

	buildInternal(*buildName, *buildPull, conf)
}

// Perform the actual build process
func buildInternal(name string, pull bool, conf string) {
	if name == "" {
		name = filepath.Base(conf)
		ext := filepath.Ext(conf)
		if ext != "" {
			name = name[:len(name)-len(ext)]
		}
	}

	config, err := ioutil.ReadFile(conf)
	if err != nil {
		log.Fatalf("Cannot open config file: %v", err)
	}

	m, err := NewConfig(config)
	if err != nil {
		log.Fatalf("Invalid config: %v", err)
	}

	containers := []*bytes.Buffer{}

	if pull {
		log.Infof("Pull kernel image: %s", m.Kernel.Image)
		err := dockerPull(m.Kernel.Image)
		if err != nil {
			log.Fatalf("Could not pull image %s: %v", m.Kernel.Image, err)
		}
	}
	// get kernel bzImage and initrd tarball from container
	// TODO examine contents to see what names they might have
	log.Infof("Extract kernel image: %s", m.Kernel.Image)
	const (
		bzimageName = "bzImage"
		ktarName    = "kernel.tar"
	)
	out, err := dockerRun(m.Kernel.Image, "tar", "cf", "-", bzimageName, ktarName)
	if err != nil {
		log.Fatalf("Failed to extract kernel image and tarball: %v", err)
	}
	buf := bytes.NewBuffer(out)
	bzimage, ktar, err := untarKernel(buf, bzimageName, ktarName)
	if err != nil {
		log.Fatalf("Could not extract bzImage and kernel filesystem from tarball. %v", err)
	}
	containers = append(containers, ktar)

	// convert init image to tarball
	if pull {
		log.Infof("Pull init: %s", m.Init)
		err := dockerPull(m.Init)
		if err != nil {
			log.Fatalf("Could not pull image %s: %v", m.Init, err)
		}
	}
	log.Infof("Process init: %s", m.Init)
	init, err := ImageExtract(m.Init, "")
	if err != nil {
		log.Fatalf("Failed to build init tarball: %v", err)
	}
	buffer := bytes.NewBuffer(init)
	containers = append(containers, buffer)

	log.Infof("Add system containers:")
	for i, image := range m.System {
		if pull {
			log.Infof("  Pull: %s", image.Image)
			err := dockerPull(image.Image)
			if err != nil {
				log.Fatalf("Could not pull image %s: %v", image.Image, err)
			}
		}
		log.Infof("  Create OCI config for %s", image.Image)
		config, err := ConfigToOCI(&image)
		if err != nil {
			log.Fatalf("Failed to run riddler to get config.json for %s: %v", image.Image, err)
		}
		so := fmt.Sprintf("%03d", i)
		path := "containers/system/" + so + "-" + image.Name
		out, err := ImageBundle(path, image.Image, config)
		if err != nil {
			log.Fatalf("Failed to extract root filesystem for %s: %v", image.Image, err)
		}
		buffer := bytes.NewBuffer(out)
		containers = append(containers, buffer)
	}

	log.Infof("Add daemon containers:")
	for _, image := range m.Daemon {
		if pull {
			log.Infof("  Pull: %s", image.Image)
			err := dockerPull(image.Image)
			if err != nil {
				log.Fatalf("Could not pull image %s: %v", image.Image, err)
			}
		}
		log.Infof("  Create OCI config for %s", image.Image)
		config, err := ConfigToOCI(&image)
		if err != nil {
			log.Fatalf("Failed to run riddler to get config.json for %s: %v", image.Image, err)
		}
		path := "containers/daemon/" + image.Name
		out, err := ImageBundle(path, image.Image, config)
		if err != nil {
			log.Fatalf("Failed to extract root filesystem for %s: %v", image.Image, err)
		}
		buffer := bytes.NewBuffer(out)
		containers = append(containers, buffer)
	}

	// add files
	buffer, err = filesystem(m)
	if err != nil {
		log.Fatalf("failed to add filesystem parts: %v", err)
	}
	containers = append(containers, buffer)

	log.Infof("Create initial ram disk")
	initrd, err := containersInitrd(containers)
	if err != nil {
		log.Fatalf("Failed to make initrd %v", err)
	}

	log.Infof("Create outputs:")
	err = outputs(m, name, bzimage.Bytes(), initrd.Bytes())
	if err != nil {
		log.Fatalf("Error writing outputs: %v", err)
	}
}

func untarKernel(buf *bytes.Buffer, bzimageName, ktarName string) (*bytes.Buffer, *bytes.Buffer, error) {
	tr := tar.NewReader(buf)

	var bzimage, ktar *bytes.Buffer

	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatalln(err)
		}
		switch hdr.Name {
		case bzimageName:
			bzimage = new(bytes.Buffer)
			_, err := io.Copy(bzimage, tr)
			if err != nil {
				return nil, nil, err
			}
		case ktarName:
			ktar = new(bytes.Buffer)
			_, err := io.Copy(bzimage, tr)
			if err != nil {
				return nil, nil, err
			}
		default:
			continue
		}
	}

	if ktar == nil || bzimage == nil {
		return nil, nil, errors.New("did not find bzImage and kernel.tar in tarball")
	}

	return bzimage, ktar, nil
}

func containersInitrd(containers []*bytes.Buffer) (*bytes.Buffer, error) {
	w := new(bytes.Buffer)
	iw := initrd.NewWriter(w)
	defer iw.Close()
	for _, file := range containers {
		_, err := initrd.Copy(iw, file)
		if err != nil {
			return nil, err
		}
	}

	return w, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2014 Ashley Jeffs
//
// 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.

package main

import (
	"context"
	"fmt"
	_ "net/http/pprof"
	"runtime/pprof"

	"flag"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/Jeffail/benthos/lib/api"
	"github.com/Jeffail/benthos/lib/buffer"
	"github.com/Jeffail/benthos/lib/input"
	"github.com/Jeffail/benthos/lib/output"
	"github.com/Jeffail/benthos/lib/processor"
	"github.com/Jeffail/benthos/lib/types"
	"github.com/Jeffail/benthos/lib/util"
	"github.com/Jeffail/benthos/lib/util/service"
	"github.com/Jeffail/benthos/lib/util/service/log"
	"github.com/Jeffail/benthos/lib/util/service/metrics"
)

//------------------------------------------------------------------------------

// Config is the benthos configuration struct.
type Config struct {
	HTTP                 api.Config       `json:"http" yaml:"http"`
	Input                input.Config     `json:"input" yaml:"input"`
	Output               output.Config    `json:"output" yaml:"output"`
	Buffer               buffer.Config    `json:"buffer" yaml:"buffer"`
	Logger               log.LoggerConfig `json:"logger" yaml:"logger"`
	Metrics              metrics.Config   `json:"metrics" yaml:"metrics"`
	SystemCloseTimeoutMS int              `json:"sys_exit_timeout_ms" yaml:"sys_exit_timeout_ms"`
}

// NewConfig returns a new configuration with default values.
func NewConfig() Config {
	metricsConf := metrics.NewConfig()
	metricsConf.Prefix = "benthos"

	return Config{
		HTTP:                 api.NewConfig(),
		Input:                input.NewConfig(),
		Output:               output.NewConfig(),
		Buffer:               buffer.NewConfig(),
		Logger:               log.NewLoggerConfig(),
		Metrics:              metricsConf,
		SystemCloseTimeoutMS: 20000,
	}
}

// Sanitised returns a sanitised copy of the Benthos configuration, meaning
// fields of no consequence (unused inputs, outputs, processors etc) are
// excluded.
func (c Config) Sanitised() (interface{}, error) {
	inConf, err := input.SanitiseConfig(c.Input)
	if err != nil {
		return nil, err
	}

	var outConf interface{}
	outConf, err = output.SanitiseConfig(c.Output)
	if err != nil {
		return nil, err
	}

	var bufConf interface{}
	bufConf, err = buffer.SanitiseConfig(c.Buffer)
	if err != nil {
		return nil, err
	}

	var metConf interface{}
	metConf, err = metrics.SanitiseConfig(c.Metrics)
	if err != nil {
		return nil, err
	}

	return struct {
		HTTP                 interface{} `json:"http" yaml:"http"`
		Input                interface{} `json:"input" yaml:"input"`
		Output               interface{} `json:"output" yaml:"output"`
		Buffer               interface{} `json:"buffer" yaml:"buffer"`
		Logger               interface{} `json:"logger" yaml:"logger"`
		Metrics              interface{} `json:"metrics" yaml:"metrics"`
		SystemCloseTimeoutMS interface{} `json:"sys_exit_timeout_ms" yaml:"sys_exit_timeout_ms"`
	}{
		HTTP:                 c.HTTP,
		Input:                inConf,
		Output:               outConf,
		Buffer:               bufConf,
		Logger:               c.Logger,
		Metrics:              metConf,
		SystemCloseTimeoutMS: c.SystemCloseTimeoutMS,
	}, nil
}

//------------------------------------------------------------------------------

// Extra flags
var (
	printInputs = flag.Bool(
		"list-inputs", false,
		"Print a list of available input options, then exit",
	)
	printOutputs = flag.Bool(
		"list-outputs", false,
		"Print a list of available output options, then exit",
	)
	printBuffers = flag.Bool(
		"list-buffers", false,
		"Print a list of available buffer options, then exit",
	)
	printProcessors = flag.Bool(
		"list-processors", false,
		"Print a list of available processor options, then exit",
	)
)

//------------------------------------------------------------------------------

// bootstrap reads cmd args and either parses and config file or prints helper
// text and exits.
func bootstrap() Config {
	config := NewConfig()

	// A list of default config paths to check for if not explicitly defined
	defaultPaths := []string{
		"/benthos.yaml",
		"/etc/benthos/config.yaml",
		"/etc/benthos.yaml",
	}

	// Override default help printing
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage: benthos [flags...]")
		fmt.Fprintln(os.Stderr, "Flags:")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr,
			"\nFor example configs use --print-yaml or --print-json\n"+
				"For a list of available inputs or outputs use --list-inputs or --list-outputs\n"+
				"For a list of available buffer options use --list-buffers\n")
	}

	// Load configuration etc
	if !service.Bootstrap(&config, defaultPaths...) {
		os.Exit(0)
	}

	// If we only want to print our inputs or outputs we should exit afterwards
	if *printInputs || *printOutputs || *printBuffers || *printProcessors {
		if *printInputs {
			fmt.Println(input.Descriptions())
		}
		if *printProcessors {
			fmt.Println(processor.Descriptions())
		}
		if *printBuffers {
			fmt.Println(buffer.Descriptions())
		}
		if *printOutputs {
			fmt.Println(output.Descriptions())
		}
		os.Exit(1)
	}

	return config
}

// createPipeline creates a pipeline based on the supplied configuration file,
// and return a closable pool of pipeline objects, a channel indicating that all
// inputs and outputs have seized, or an error.
func createPipeline(
	config Config, mgr types.Manager, logger log.Modular, stats metrics.Type,
) (*util.ClosablePool, *util.ClosablePool, chan struct{}, error) {
	// Create two pools, this helps manage ordered closure of all pipeline
	// components. We have a tiered (t1) and an non-tiered (t2) pool. If the
	// tiered pool cannot close within our allotted time period then we try
	// closing the second non-tiered pool. If the second pool also fails then we
	// exit the service ungracefully.
	poolt1, poolt2 := util.NewClosablePool(), util.NewClosablePool()

	// Create our input pipe
	inputPipe, err := input.New(config.Input, mgr, logger, stats)
	if err != nil {
		logger.Errorf("Input error (%s): %v\n", config.Input.Type, err)
		return nil, nil, nil, err
	}
	poolt1.Add(1, inputPipe)
	poolt2.Add(0, inputPipe)

	// Create a buffer
	buf, err := buffer.New(config.Buffer, logger, stats)
	if err != nil {
		logger.Errorf("Buffer error (%s): %v\n", config.Buffer.Type, err)
		return nil, nil, nil, err
	}
	poolt1.Add(3, buf)
	poolt2.Add(0, buf)

	// Create our output pipe
	outputPipe, err := output.New(config.Output, mgr, logger, stats)
	if err != nil {
		logger.Errorf("Output error (%s): %v\n", config.Output.Type, err)
		return nil, nil, nil, err
	}
	poolt1.Add(10, outputPipe)
	poolt2.Add(0, outputPipe)

	util.Couple(buf, outputPipe)
	util.Couple(inputPipe, buf)
	closeChan := make(chan struct{})

	// If our outputs close down then we should shut down the service
	go func() {
		for {
			if err := outputPipe.WaitForClose(time.Second * 60); err == nil {
				closeChan <- struct{}{}
				return
			}
		}
	}()

	return poolt1, poolt2, closeChan, nil
}

func main() {
	// Bootstrap by reading cmd flags and configuration file
	config := bootstrap()

	// Logging and stats aggregation
	var logger log.Modular

	// Note: Only log to Stderr if one of our outputs is stdout
	if config.Output.Type == "stdout" {
		logger = log.NewLogger(os.Stderr, config.Logger)
	} else {
		logger = log.NewLogger(os.Stdout, config.Logger)
	}

	logger.Infoln("Launching a benthos instance, use CTRL+C to close.")

	// Create our metrics type
	stats, err := metrics.New(config.Metrics)
	if err != nil {
		logger.Errorf("Metrics error: %v\n", err)
		return
	}
	defer stats.Close()

	sanConf, err := config.Sanitised()
	if err != nil {
		logger.Warnf("Failed to generate sanitised config: %v\n", err)
	}
	httpServer := api.New(service.Version, service.DateBuilt, config.HTTP, sanConf, logger, stats)

	poolTiered, poolNonTiered, outputsClosedChan, err := createPipeline(config, httpServer, logger, stats)
	if err != nil {
		logger.Errorf("Service closing due to: %v\n", err)
		return
	}

	httpServerClosedChan := make(chan struct{})
	go func() {
		logger.Infof(
			"Listening for HTTP requests at: %v\n",
			"http://"+config.HTTP.Address,
		)
		httpErr := httpServer.ListenAndServe()
		if httpErr != nil && httpErr != http.ErrServerClosed {
			logger.Errorf("HTTP Server error: %v\n", httpErr)
		}
		close(httpServerClosedChan)
	}()

	// Defer ordered pool clean up.
	defer func() {
		tout := time.Millisecond * time.Duration(config.SystemCloseTimeoutMS)

		go func() {
			httpServer.Shutdown(context.Background())
			select {
			case <-httpServerClosedChan:
			case <-time.After(tout / 2):
				logger.Warnln("Service failed to close HTTP server gracefully in time.")
			}
		}()

		if config.Logger.LogLevel == "DEBUG" {
			go func() {
				<-time.After(tout)
				pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
				os.Exit(1)
			}()
		}

		if err := poolTiered.Close(tout / 2); err != nil {
			logger.Warnln(
				"Service failed to close using ordered tiers, you may receive a duplicate " +
					"message on the next service start.",
			)
			if err = poolNonTiered.Close(tout / 2); err != nil {
				logger.Warnln(
					"Service failed to close cleanly within allocated time. Exiting forcefully.",
				)
				os.Exit(1)
			}
		}
	}()

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

	// Wait for termination signal
	select {
	case <-sigChan:
		logger.Infoln("Received SIGTERM, the service is closing.")
	case <-outputsClosedChan:
		logger.Infoln("Pipeline has terminated. Shutting down the service.")
	case <-httpServerClosedChan:
		logger.Infoln("HTTP Server has terminated. Shutting down the service.")
	}
}

//------------------------------------------------------------------------------
<commit_msg>Guarantee shutdown<commit_after>// Copyright (c) 2014 Ashley Jeffs
//
// 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.

package main

import (
	"context"
	"fmt"
	_ "net/http/pprof"
	"runtime/pprof"

	"flag"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/Jeffail/benthos/lib/api"
	"github.com/Jeffail/benthos/lib/buffer"
	"github.com/Jeffail/benthos/lib/input"
	"github.com/Jeffail/benthos/lib/output"
	"github.com/Jeffail/benthos/lib/processor"
	"github.com/Jeffail/benthos/lib/types"
	"github.com/Jeffail/benthos/lib/util"
	"github.com/Jeffail/benthos/lib/util/service"
	"github.com/Jeffail/benthos/lib/util/service/log"
	"github.com/Jeffail/benthos/lib/util/service/metrics"
)

//------------------------------------------------------------------------------

// Config is the benthos configuration struct.
type Config struct {
	HTTP                 api.Config       `json:"http" yaml:"http"`
	Input                input.Config     `json:"input" yaml:"input"`
	Output               output.Config    `json:"output" yaml:"output"`
	Buffer               buffer.Config    `json:"buffer" yaml:"buffer"`
	Logger               log.LoggerConfig `json:"logger" yaml:"logger"`
	Metrics              metrics.Config   `json:"metrics" yaml:"metrics"`
	SystemCloseTimeoutMS int              `json:"sys_exit_timeout_ms" yaml:"sys_exit_timeout_ms"`
}

// NewConfig returns a new configuration with default values.
func NewConfig() Config {
	metricsConf := metrics.NewConfig()
	metricsConf.Prefix = "benthos"

	return Config{
		HTTP:                 api.NewConfig(),
		Input:                input.NewConfig(),
		Output:               output.NewConfig(),
		Buffer:               buffer.NewConfig(),
		Logger:               log.NewLoggerConfig(),
		Metrics:              metricsConf,
		SystemCloseTimeoutMS: 20000,
	}
}

// Sanitised returns a sanitised copy of the Benthos configuration, meaning
// fields of no consequence (unused inputs, outputs, processors etc) are
// excluded.
func (c Config) Sanitised() (interface{}, error) {
	inConf, err := input.SanitiseConfig(c.Input)
	if err != nil {
		return nil, err
	}

	var outConf interface{}
	outConf, err = output.SanitiseConfig(c.Output)
	if err != nil {
		return nil, err
	}

	var bufConf interface{}
	bufConf, err = buffer.SanitiseConfig(c.Buffer)
	if err != nil {
		return nil, err
	}

	var metConf interface{}
	metConf, err = metrics.SanitiseConfig(c.Metrics)
	if err != nil {
		return nil, err
	}

	return struct {
		HTTP                 interface{} `json:"http" yaml:"http"`
		Input                interface{} `json:"input" yaml:"input"`
		Output               interface{} `json:"output" yaml:"output"`
		Buffer               interface{} `json:"buffer" yaml:"buffer"`
		Logger               interface{} `json:"logger" yaml:"logger"`
		Metrics              interface{} `json:"metrics" yaml:"metrics"`
		SystemCloseTimeoutMS interface{} `json:"sys_exit_timeout_ms" yaml:"sys_exit_timeout_ms"`
	}{
		HTTP:                 c.HTTP,
		Input:                inConf,
		Output:               outConf,
		Buffer:               bufConf,
		Logger:               c.Logger,
		Metrics:              metConf,
		SystemCloseTimeoutMS: c.SystemCloseTimeoutMS,
	}, nil
}

//------------------------------------------------------------------------------

// Extra flags
var (
	printInputs = flag.Bool(
		"list-inputs", false,
		"Print a list of available input options, then exit",
	)
	printOutputs = flag.Bool(
		"list-outputs", false,
		"Print a list of available output options, then exit",
	)
	printBuffers = flag.Bool(
		"list-buffers", false,
		"Print a list of available buffer options, then exit",
	)
	printProcessors = flag.Bool(
		"list-processors", false,
		"Print a list of available processor options, then exit",
	)
)

//------------------------------------------------------------------------------

// bootstrap reads cmd args and either parses and config file or prints helper
// text and exits.
func bootstrap() Config {
	config := NewConfig()

	// A list of default config paths to check for if not explicitly defined
	defaultPaths := []string{
		"/benthos.yaml",
		"/etc/benthos/config.yaml",
		"/etc/benthos.yaml",
	}

	// Override default help printing
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage: benthos [flags...]")
		fmt.Fprintln(os.Stderr, "Flags:")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr,
			"\nFor example configs use --print-yaml or --print-json\n"+
				"For a list of available inputs or outputs use --list-inputs or --list-outputs\n"+
				"For a list of available buffer options use --list-buffers\n")
	}

	// Load configuration etc
	if !service.Bootstrap(&config, defaultPaths...) {
		os.Exit(0)
	}

	// If we only want to print our inputs or outputs we should exit afterwards
	if *printInputs || *printOutputs || *printBuffers || *printProcessors {
		if *printInputs {
			fmt.Println(input.Descriptions())
		}
		if *printProcessors {
			fmt.Println(processor.Descriptions())
		}
		if *printBuffers {
			fmt.Println(buffer.Descriptions())
		}
		if *printOutputs {
			fmt.Println(output.Descriptions())
		}
		os.Exit(1)
	}

	return config
}

// createPipeline creates a pipeline based on the supplied configuration file,
// and return a closable pool of pipeline objects, a channel indicating that all
// inputs and outputs have seized, or an error.
func createPipeline(
	config Config, mgr types.Manager, logger log.Modular, stats metrics.Type,
) (*util.ClosablePool, *util.ClosablePool, chan struct{}, error) {
	// Create two pools, this helps manage ordered closure of all pipeline
	// components. We have a tiered (t1) and an non-tiered (t2) pool. If the
	// tiered pool cannot close within our allotted time period then we try
	// closing the second non-tiered pool. If the second pool also fails then we
	// exit the service ungracefully.
	poolt1, poolt2 := util.NewClosablePool(), util.NewClosablePool()

	// Create our input pipe
	inputPipe, err := input.New(config.Input, mgr, logger, stats)
	if err != nil {
		logger.Errorf("Input error (%s): %v\n", config.Input.Type, err)
		return nil, nil, nil, err
	}
	poolt1.Add(1, inputPipe)
	poolt2.Add(0, inputPipe)

	// Create a buffer
	buf, err := buffer.New(config.Buffer, logger, stats)
	if err != nil {
		logger.Errorf("Buffer error (%s): %v\n", config.Buffer.Type, err)
		return nil, nil, nil, err
	}
	poolt1.Add(3, buf)
	poolt2.Add(0, buf)

	// Create our output pipe
	outputPipe, err := output.New(config.Output, mgr, logger, stats)
	if err != nil {
		logger.Errorf("Output error (%s): %v\n", config.Output.Type, err)
		return nil, nil, nil, err
	}
	poolt1.Add(10, outputPipe)
	poolt2.Add(0, outputPipe)

	util.Couple(buf, outputPipe)
	util.Couple(inputPipe, buf)
	closeChan := make(chan struct{})

	// If our outputs close down then we should shut down the service
	go func() {
		for {
			if err := outputPipe.WaitForClose(time.Second * 60); err == nil {
				closeChan <- struct{}{}
				return
			}
		}
	}()

	return poolt1, poolt2, closeChan, nil
}

func main() {
	// Bootstrap by reading cmd flags and configuration file
	config := bootstrap()

	// Logging and stats aggregation
	var logger log.Modular

	// Note: Only log to Stderr if one of our outputs is stdout
	if config.Output.Type == "stdout" {
		logger = log.NewLogger(os.Stderr, config.Logger)
	} else {
		logger = log.NewLogger(os.Stdout, config.Logger)
	}

	logger.Infoln("Launching a benthos instance, use CTRL+C to close.")

	// Create our metrics type
	stats, err := metrics.New(config.Metrics)
	if err != nil {
		logger.Errorf("Metrics error: %v\n", err)
		return
	}
	defer stats.Close()

	sanConf, err := config.Sanitised()
	if err != nil {
		logger.Warnf("Failed to generate sanitised config: %v\n", err)
	}
	httpServer := api.New(service.Version, service.DateBuilt, config.HTTP, sanConf, logger, stats)

	poolTiered, poolNonTiered, outputsClosedChan, err := createPipeline(config, httpServer, logger, stats)
	if err != nil {
		logger.Errorf("Service closing due to: %v\n", err)
		return
	}

	httpServerClosedChan := make(chan struct{})
	go func() {
		logger.Infof(
			"Listening for HTTP requests at: %v\n",
			"http://"+config.HTTP.Address,
		)
		httpErr := httpServer.ListenAndServe()
		if httpErr != nil && httpErr != http.ErrServerClosed {
			logger.Errorf("HTTP Server error: %v\n", httpErr)
		}
		close(httpServerClosedChan)
	}()

	// Defer ordered pool clean up.
	defer func() {
		tout := time.Millisecond * time.Duration(config.SystemCloseTimeoutMS)

		go func() {
			httpServer.Shutdown(context.Background())
			select {
			case <-httpServerClosedChan:
			case <-time.After(tout / 2):
				logger.Warnln("Service failed to close HTTP server gracefully in time.")
			}
		}()

		go func() {
			<-time.After(tout + time.Second)
			logger.Warnln(
				"Service failed to close cleanly within allocated time. Exiting forcefully.",
			)
			if config.Logger.LogLevel == "DEBUG" {
				pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
			}
			os.Exit(1)
		}()

		if err := poolTiered.Close(tout / 2); err != nil {
			logger.Warnln(
				"Service failed to close using ordered tiers, you may receive a duplicate " +
					"message on the next service start.",
			)
			if err = poolNonTiered.Close(tout / 2); err != nil {
				logger.Warnln(
					"Service failed to close cleanly within allocated time. Exiting forcefully.",
				)
				os.Exit(1)
			}
		}
	}()

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

	// Wait for termination signal
	select {
	case <-sigChan:
		logger.Infoln("Received SIGTERM, the service is closing.")
	case <-outputsClosedChan:
		logger.Infoln("Pipeline has terminated. Shutting down the service.")
	case <-httpServerClosedChan:
		logger.Infoln("HTTP Server has terminated. Shutting down the service.")
	}
}

//------------------------------------------------------------------------------
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"math"

	"github.com/ready-steady/format/mat"
	"github.com/ready-steady/numeric/interpolation/adhier"
	"github.com/ready-steady/statistics"
	"github.com/ready-steady/statistics/test"

	"../internal"
)

const (
	deltaCensiusKelvin = 273.15
)

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

func command(config internal.Config, input *mat.File, _ *mat.File) error {
	if input == nil {
		return errors.New("an input file is required")
	}

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

	observations := []float64{}
	if err := input.Get("observations", &observations); err != nil {
		return err
	}

	predictions := []float64{}
	if err := input.Get("predictions", &predictions); err != nil {
		return err
	}

	ns := int(config.Assessment.Samples)
	no := len(observations) / ns

	cut := func(data []float64, i int) []float64 {
		piece := make([]float64, ns)
		for j := 0; j < ns; j++ {
			piece[j] = data[j*no+i]
		}
		return piece
	}

	fmt.Printf("Surrogate: inputs %d, outputs %d, level %d, nodes %d\n",
		surrogate.Inputs, surrogate.Outputs, surrogate.Level, surrogate.Nodes)

	εμ := make([]float64, no)
	εσ := make([]float64, no)
	εp := make([]float64, no)

	// Compute errors across all outputs.
	for i := 0; i < no; i++ {
		observations := cut(observations, i)
		predictions := cut(predictions, i)

		μ1 := statistics.Mean(observations)
		μ2 := statistics.Mean(predictions)
		εμ[i] = math.Abs(μ1 - μ2)

		σ1 := math.Sqrt(statistics.Variance(observations))
		σ2 := math.Sqrt(statistics.Variance(predictions))
		εσ[i] = math.Abs(σ1 - σ2)

		_, _, εp[i] = test.KolmogorovSmirnov(observations, predictions, 0)

		if no > 1 && config.Verbose {
			fmt.Printf("%9d: μ %10.2e ±%10.2e, σ %10.2e ±%10.2e, p %.2e\n",
				i, μ1-deltaCensiusKelvin, εμ[i], σ1, εσ[i], εp[i])
		}
	}

	if no == 1 {
		fmt.Printf("Error: μ ±%.2e, σ ±%.2e, p %.2e\n", εμ[0], εσ[0], εp[0])
	} else {
		fmt.Printf("Average error: μ ±%10.2e, σ ±%10.2e, p %.2e\n",
			statistics.Mean(εμ), statistics.Mean(εσ), statistics.Mean(εp))

		fmt.Printf("Maximal error: μ ±%10.2e, σ ±%10.2e, p %.2e\n",
			max(εμ), max(εσ), max(εp))
	}

	return nil
}

func max(data []float64) float64 {
	max := math.Inf(-1)

	for _, x := range data {
		if x > max {
			max = x
		}
	}

	return max
}
<commit_msg>Adjusted the output of cmd/compare<commit_after>package main

import (
	"errors"
	"fmt"
	"math"

	"github.com/ready-steady/format/mat"
	"github.com/ready-steady/numeric/interpolation/adhier"
	"github.com/ready-steady/statistics"
	"github.com/ready-steady/statistics/test"

	"../internal"
)

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

func command(config internal.Config, input *mat.File, _ *mat.File) error {
	if input == nil {
		return errors.New("an input file is required")
	}

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

	observations := []float64{}
	if err := input.Get("observations", &observations); err != nil {
		return err
	}

	predictions := []float64{}
	if err := input.Get("predictions", &predictions); err != nil {
		return err
	}

	ns := int(config.Assessment.Samples)
	no := len(observations) / ns

	cut := func(data []float64, i int) []float64 {
		piece := make([]float64, ns)
		for j := 0; j < ns; j++ {
			piece[j] = data[j*no+i]
		}
		return piece
	}

	fmt.Printf("Surrogate: inputs %d, outputs %d, level %d, nodes %d\n",
		surrogate.Inputs, surrogate.Outputs, surrogate.Level, surrogate.Nodes)

	εμ := make([]float64, no)
	εσ := make([]float64, no)
	εp := make([]float64, no)

	// Compute errors across all outputs.
	for i := 0; i < no; i++ {
		observations := cut(observations, i)
		predictions := cut(predictions, i)

		μ1 := statistics.Mean(observations)
		μ2 := statistics.Mean(predictions)
		εμ[i] = math.Abs(μ1 - μ2)

		σ1 := math.Sqrt(statistics.Variance(observations))
		σ2 := math.Sqrt(statistics.Variance(predictions))
		εσ[i] = math.Abs(σ1 - σ2)

		_, _, εp[i] = test.KolmogorovSmirnov(observations, predictions, 0)

		if no == 1 {
			fmt.Printf("Error: μ %10.2e ±%10.2e, σ %10.2e ±%10.2e, p %.2e\n",
				μ1, εμ[i], σ1, εσ[i], εp[i])
		} else if config.Verbose {
			fmt.Printf("%9d: μ %10.2e ±%10.2e, σ %10.2e ±%10.2e, p %.2e\n",
				i, μ1, εμ[i], σ1, εσ[i], εp[i])
		}
	}

	if no > 1 {
		fmt.Printf("Average error: μ ±%10.2e, σ ±%10.2e, p %.2e\n",
			statistics.Mean(εμ), statistics.Mean(εσ), statistics.Mean(εp))

		fmt.Printf("Maximal error: μ ±%10.2e, σ ±%10.2e, p %.2e\n",
			max(εμ), max(εσ), max(εp))
	}

	return nil
}

func max(data []float64) float64 {
	max := math.Inf(-1)

	for _, x := range data {
		if x > max {
			max = x
		}
	}

	return max
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"github.com/jkomoros/sudoku"
	"github.com/jkomoros/sudoku/sdkconverter"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
)

//TODO: let people pass in a filename to export to.

const STORED_PUZZLES_DIRECTORY = ".puzzles"

type appOptions struct {
	GENERATE            bool
	HELP                bool
	PUZZLE_TO_SOLVE     string
	NUM                 int
	PRINT_STATS         bool
	WALKTHROUGH         bool
	RAW_SYMMETRY        string
	SYMMETRY            sudoku.SymmetryType
	SYMMETRY_PROPORTION float64
	MIN_DIFFICULTY      float64
	MAX_DIFFICULTY      float64
	NO_CACHE            bool
	PUZZLE_FORMAT       string
	OUTPUT_CSV          bool
	CONVERTER           sdkconverter.SudokuPuzzleConverter
}

func init() {
	//grid.Difficulty can make use of a number of processes simultaneously.
	runtime.GOMAXPROCS(6)
}

func main() {

	//TODO: figure out how to test this.

	var options appOptions

	flag.BoolVar(&options.GENERATE, "g", false, "if true, will generate a puzzle.")
	flag.BoolVar(&options.HELP, "h", false, "If provided, will print help and exit.")
	flag.IntVar(&options.NUM, "n", 1, "Number of things to generate")
	flag.BoolVar(&options.PRINT_STATS, "p", false, "If provided, will print stats.")
	flag.StringVar(&options.PUZZLE_TO_SOLVE, "s", "", "If provided, will solve the puzzle at the given filename and print solution.")
	flag.BoolVar(&options.WALKTHROUGH, "w", false, "If provided, will print out a walkthrough to solve the provided puzzle.")
	flag.StringVar(&options.RAW_SYMMETRY, "y", "vertical", "Valid values: 'none', 'both', 'horizontal', 'vertical")
	flag.Float64Var(&options.SYMMETRY_PROPORTION, "r", 0.7, "What proportion of cells should be filled according to symmetry")
	flag.Float64Var(&options.MIN_DIFFICULTY, "min", 0.0, "Minimum difficulty for generated puzzle")
	flag.Float64Var(&options.MAX_DIFFICULTY, "max", 1.0, "Maximum difficulty for generated puzzle")
	flag.BoolVar(&options.NO_CACHE, "no-cache", false, "If provided, will not vend generated puzzles from the cache of previously generated puzzles.")
	//TODO: the format should also be how we interpret loads, too.
	flag.StringVar(&options.PUZZLE_FORMAT, "format", "sdk", "Which format to export puzzles from. Defaults to 'sdk'")
	flag.BoolVar(&options.OUTPUT_CSV, "csv", false, "Output the results in CSV.")
	flag.Parse()

	options.RAW_SYMMETRY = strings.ToLower(options.RAW_SYMMETRY)
	switch options.RAW_SYMMETRY {
	case "none":
		options.SYMMETRY = sudoku.SYMMETRY_NONE
	case "both":
		options.SYMMETRY = sudoku.SYMMETRY_BOTH
	case "horizontal":
		options.SYMMETRY = sudoku.SYMMETRY_HORIZONTAL
	case "vertical":
		options.SYMMETRY = sudoku.SYMMETRY_VERTICAL
	default:
		log.Fatal("Unknown symmetry flag: ", options.RAW_SYMMETRY)
	}

	options.CONVERTER = sdkconverter.Converters[options.PUZZLE_FORMAT]

	if options.CONVERTER == nil {
		log.Fatal("Invalid format option:", options.PUZZLE_FORMAT)
	}

	output := os.Stdout

	if options.HELP {
		flag.PrintDefaults()
		return
	}

	var grid *sudoku.Grid

	var csvWriter *csv.Writer
	var csvRec []string

	if options.OUTPUT_CSV {
		csvWriter = csv.NewWriter(output)
	}

	for i := 0; i < options.NUM; i++ {

		if options.OUTPUT_CSV {
			csvRec = nil
		}

		//TODO: allow the type of symmetry to be configured.
		if options.GENERATE {
			grid = generatePuzzle(options.MIN_DIFFICULTY, options.MAX_DIFFICULTY, options.SYMMETRY, options.SYMMETRY_PROPORTION, options.NO_CACHE)
			//TODO: factor out all of this double-printing.
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, options.CONVERTER.DataString(grid))
			} else {
				fmt.Fprintln(output, options.CONVERTER.DataString(grid))
			}
		} else if options.PUZZLE_TO_SOLVE != "" {
			//TODO: detect if the load failed.
			grid = sudoku.NewGrid()

			data, err := ioutil.ReadFile(options.PUZZLE_TO_SOLVE)

			if err != nil {
				log.Fatalln("Read error for specified file:", err)
			}

			//TODO: shouldn't a load method have a way to say the string provided is invalid?
			options.CONVERTER.Load(grid, string(data))
		}

		if grid == nil {
			//No grid to do anything with.
			log.Fatalln("No grid loaded.")
		}

		//TODO: use of this option leads to a busy loop somewhere... Is it related to the generate-multiple-and-difficulty hang?

		var directions sudoku.SolveDirections

		if options.WALKTHROUGH || options.PRINT_STATS {
			directions = grid.HumanSolution()
			if len(directions) == 0 {
				//We couldn't solve it. Let's check and see if the puzzle is well formed.
				if grid.HasMultipleSolutions() {
					//TODO: figure out why guesses wouldn't be used here effectively.
					log.Println("The puzzle had multiple solutions; that means it's not well-formed")
				}
			}
		}

		if options.WALKTHROUGH {
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, directions.Walkthrough(grid))
			} else {
				fmt.Fprintln(output, directions.Walkthrough(grid))
			}
		}
		if options.PRINT_STATS {
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, strconv.FormatFloat(grid.Difficulty(), 'f', -1, 64))
				//We won't print out the directions.Stats() like we do for just printing to stdout,
				//because that's mostly noise in this format.
			} else {
				fmt.Fprintln(output, grid.Difficulty())
				//TODO: consider actually printing out the Signals stats (with a Stats method on signals)
				fmt.Fprintln(output, strings.Join(directions.Stats(), "\n"))
			}
		}
		if options.PUZZLE_TO_SOLVE != "" {
			grid.Solve()
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, options.CONVERTER.DataString(grid))
			} else {
				fmt.Fprintln(output, options.CONVERTER.DataString(grid))

			}
		}

		if options.OUTPUT_CSV {
			csvWriter.Write(csvRec)
		}

		if options.PUZZLE_TO_SOLVE != "" {
			//If we're asked to solve, n could only be 1 anyway.
			return
		}
		grid.Done()
	}
	if options.OUTPUT_CSV {
		csvWriter.Flush()
	}

}

func puzzleDirectoryParts(symmetryType sudoku.SymmetryType, symmetryPercentage float64) []string {
	return []string{
		STORED_PUZZLES_DIRECTORY,
		"SYM_TYPE_" + strconv.Itoa(int(symmetryType)),
		"SYM_PERCENTAGE_" + strconv.FormatFloat(symmetryPercentage, 'f', -1, 64),
	}
}

func storePuzzle(grid *sudoku.Grid, difficulty float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64) bool {
	//TODO: we should include a hashed version of our difficulty weights file so we don't cache ones with old weights.
	directoryParts := puzzleDirectoryParts(symmetryType, symmetryPercentage)

	fileNamePart := strconv.FormatFloat(difficulty, 'f', -1, 64) + ".sdk"

	pathSoFar := ""

	for i, part := range directoryParts {
		if i == 0 {
			pathSoFar = part
		} else {
			pathSoFar = filepath.Join(pathSoFar, part)
		}
		if _, err := os.Stat(pathSoFar); os.IsNotExist(err) {
			//need to create it.
			os.Mkdir(pathSoFar, 0700)
		}
	}

	fileName := filepath.Join(pathSoFar, fileNamePart)

	file, err := os.Create(fileName)

	if err != nil {
		log.Println(err)
		return false
	}

	defer file.Close()

	puzzleText := grid.DataString()

	n, err := io.WriteString(file, puzzleText)

	if err != nil {
		log.Println(err)
		return false
	} else {
		if n < len(puzzleText) {
			log.Println("Didn't write full file, only wrote", n, "bytes of", len(puzzleText))
			return false
		}
	}
	return true
}

func vendPuzzle(min float64, max float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64) *sudoku.Grid {

	directory := filepath.Join(puzzleDirectoryParts(symmetryType, symmetryPercentage)...)

	if files, err := ioutil.ReadDir(directory); os.IsNotExist(err) {
		//The directory doesn't exist.
		return nil
	} else {
		//OK, the directory exists, now see which puzzles are there and if any fit. If one does, vend it and delete the file.
		for _, file := range files {
			//See what this actually returns.
			filenameParts := strings.Split(file.Name(), ".")

			//Remember: there's a dot in the filename due to the float seperator.
			//TODO: shouldn't "sdk" be in a constant somewhere?
			if len(filenameParts) != 3 || filenameParts[2] != "sdk" {
				continue
			}

			difficulty, err := strconv.ParseFloat(strings.Join(filenameParts[0:2], "."), 64)
			if err != nil {
				continue
			}

			if min <= difficulty && difficulty <= max {
				//Found a puzzle!
				grid := sudoku.NewGrid()
				fullFileName := filepath.Join(directory, file.Name())
				grid.LoadFromFile(fullFileName)
				os.Remove(fullFileName)
				return grid
			}
		}
	}
	return nil
}

func generatePuzzle(min float64, max float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64, skipCache bool) *sudoku.Grid {
	var result *sudoku.Grid

	if !skipCache {
		result = vendPuzzle(min, max, symmetryType, symmetryPercentage)

		if result != nil {
			log.Println("Vending a puzzle from the cache.")
			return result
		}
	}

	//We'll have to generate one ourselves.
	count := 0
	for {
		//The first time we don't bother saying what number attemp it is, because if the first run is likely to generate a useable puzzle it's just noise.
		if count != 0 {
			log.Println("Attempt", count, "at generating puzzle.")
		}

		result = sudoku.GenerateGrid(symmetryType, symmetryPercentage)

		difficulty := result.Difficulty()

		if difficulty >= min && difficulty <= max {
			return result
		}

		log.Println("Rejecting grid of difficulty", difficulty)
		if storePuzzle(result, difficulty, symmetryType, symmetryPercentage) {
			log.Println("Stored the puzzle for future use.")
		}

		count++
	}
	return nil
}
<commit_msg>GoDoc for dokugen<commit_after>/*
dokugen is a simple command line utility that exposes many of the basic functions of the
sudoku package. It's able to generate puzzles (with difficutly) and solve provided puzzles.
Run with -h to see help on how to use it.
*/
package main

import (
	"encoding/csv"
	"flag"
	"fmt"
	"github.com/jkomoros/sudoku"
	"github.com/jkomoros/sudoku/sdkconverter"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
)

//TODO: let people pass in a filename to export to.

const STORED_PUZZLES_DIRECTORY = ".puzzles"

type appOptions struct {
	GENERATE            bool
	HELP                bool
	PUZZLE_TO_SOLVE     string
	NUM                 int
	PRINT_STATS         bool
	WALKTHROUGH         bool
	RAW_SYMMETRY        string
	SYMMETRY            sudoku.SymmetryType
	SYMMETRY_PROPORTION float64
	MIN_DIFFICULTY      float64
	MAX_DIFFICULTY      float64
	NO_CACHE            bool
	PUZZLE_FORMAT       string
	OUTPUT_CSV          bool
	CONVERTER           sdkconverter.SudokuPuzzleConverter
}

func init() {
	//grid.Difficulty can make use of a number of processes simultaneously.
	runtime.GOMAXPROCS(6)
}

func main() {

	//TODO: figure out how to test this.

	var options appOptions

	flag.BoolVar(&options.GENERATE, "g", false, "if true, will generate a puzzle.")
	flag.BoolVar(&options.HELP, "h", false, "If provided, will print help and exit.")
	flag.IntVar(&options.NUM, "n", 1, "Number of things to generate")
	flag.BoolVar(&options.PRINT_STATS, "p", false, "If provided, will print stats.")
	flag.StringVar(&options.PUZZLE_TO_SOLVE, "s", "", "If provided, will solve the puzzle at the given filename and print solution.")
	flag.BoolVar(&options.WALKTHROUGH, "w", false, "If provided, will print out a walkthrough to solve the provided puzzle.")
	flag.StringVar(&options.RAW_SYMMETRY, "y", "vertical", "Valid values: 'none', 'both', 'horizontal', 'vertical")
	flag.Float64Var(&options.SYMMETRY_PROPORTION, "r", 0.7, "What proportion of cells should be filled according to symmetry")
	flag.Float64Var(&options.MIN_DIFFICULTY, "min", 0.0, "Minimum difficulty for generated puzzle")
	flag.Float64Var(&options.MAX_DIFFICULTY, "max", 1.0, "Maximum difficulty for generated puzzle")
	flag.BoolVar(&options.NO_CACHE, "no-cache", false, "If provided, will not vend generated puzzles from the cache of previously generated puzzles.")
	//TODO: the format should also be how we interpret loads, too.
	flag.StringVar(&options.PUZZLE_FORMAT, "format", "sdk", "Which format to export puzzles from. Defaults to 'sdk'")
	flag.BoolVar(&options.OUTPUT_CSV, "csv", false, "Output the results in CSV.")
	flag.Parse()

	options.RAW_SYMMETRY = strings.ToLower(options.RAW_SYMMETRY)
	switch options.RAW_SYMMETRY {
	case "none":
		options.SYMMETRY = sudoku.SYMMETRY_NONE
	case "both":
		options.SYMMETRY = sudoku.SYMMETRY_BOTH
	case "horizontal":
		options.SYMMETRY = sudoku.SYMMETRY_HORIZONTAL
	case "vertical":
		options.SYMMETRY = sudoku.SYMMETRY_VERTICAL
	default:
		log.Fatal("Unknown symmetry flag: ", options.RAW_SYMMETRY)
	}

	options.CONVERTER = sdkconverter.Converters[options.PUZZLE_FORMAT]

	if options.CONVERTER == nil {
		log.Fatal("Invalid format option:", options.PUZZLE_FORMAT)
	}

	output := os.Stdout

	if options.HELP {
		flag.PrintDefaults()
		return
	}

	var grid *sudoku.Grid

	var csvWriter *csv.Writer
	var csvRec []string

	if options.OUTPUT_CSV {
		csvWriter = csv.NewWriter(output)
	}

	for i := 0; i < options.NUM; i++ {

		if options.OUTPUT_CSV {
			csvRec = nil
		}

		//TODO: allow the type of symmetry to be configured.
		if options.GENERATE {
			grid = generatePuzzle(options.MIN_DIFFICULTY, options.MAX_DIFFICULTY, options.SYMMETRY, options.SYMMETRY_PROPORTION, options.NO_CACHE)
			//TODO: factor out all of this double-printing.
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, options.CONVERTER.DataString(grid))
			} else {
				fmt.Fprintln(output, options.CONVERTER.DataString(grid))
			}
		} else if options.PUZZLE_TO_SOLVE != "" {
			//TODO: detect if the load failed.
			grid = sudoku.NewGrid()

			data, err := ioutil.ReadFile(options.PUZZLE_TO_SOLVE)

			if err != nil {
				log.Fatalln("Read error for specified file:", err)
			}

			//TODO: shouldn't a load method have a way to say the string provided is invalid?
			options.CONVERTER.Load(grid, string(data))
		}

		if grid == nil {
			//No grid to do anything with.
			log.Fatalln("No grid loaded.")
		}

		//TODO: use of this option leads to a busy loop somewhere... Is it related to the generate-multiple-and-difficulty hang?

		var directions sudoku.SolveDirections

		if options.WALKTHROUGH || options.PRINT_STATS {
			directions = grid.HumanSolution()
			if len(directions) == 0 {
				//We couldn't solve it. Let's check and see if the puzzle is well formed.
				if grid.HasMultipleSolutions() {
					//TODO: figure out why guesses wouldn't be used here effectively.
					log.Println("The puzzle had multiple solutions; that means it's not well-formed")
				}
			}
		}

		if options.WALKTHROUGH {
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, directions.Walkthrough(grid))
			} else {
				fmt.Fprintln(output, directions.Walkthrough(grid))
			}
		}
		if options.PRINT_STATS {
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, strconv.FormatFloat(grid.Difficulty(), 'f', -1, 64))
				//We won't print out the directions.Stats() like we do for just printing to stdout,
				//because that's mostly noise in this format.
			} else {
				fmt.Fprintln(output, grid.Difficulty())
				//TODO: consider actually printing out the Signals stats (with a Stats method on signals)
				fmt.Fprintln(output, strings.Join(directions.Stats(), "\n"))
			}
		}
		if options.PUZZLE_TO_SOLVE != "" {
			grid.Solve()
			if options.OUTPUT_CSV {
				csvRec = append(csvRec, options.CONVERTER.DataString(grid))
			} else {
				fmt.Fprintln(output, options.CONVERTER.DataString(grid))

			}
		}

		if options.OUTPUT_CSV {
			csvWriter.Write(csvRec)
		}

		if options.PUZZLE_TO_SOLVE != "" {
			//If we're asked to solve, n could only be 1 anyway.
			return
		}
		grid.Done()
	}
	if options.OUTPUT_CSV {
		csvWriter.Flush()
	}

}

func puzzleDirectoryParts(symmetryType sudoku.SymmetryType, symmetryPercentage float64) []string {
	return []string{
		STORED_PUZZLES_DIRECTORY,
		"SYM_TYPE_" + strconv.Itoa(int(symmetryType)),
		"SYM_PERCENTAGE_" + strconv.FormatFloat(symmetryPercentage, 'f', -1, 64),
	}
}

func storePuzzle(grid *sudoku.Grid, difficulty float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64) bool {
	//TODO: we should include a hashed version of our difficulty weights file so we don't cache ones with old weights.
	directoryParts := puzzleDirectoryParts(symmetryType, symmetryPercentage)

	fileNamePart := strconv.FormatFloat(difficulty, 'f', -1, 64) + ".sdk"

	pathSoFar := ""

	for i, part := range directoryParts {
		if i == 0 {
			pathSoFar = part
		} else {
			pathSoFar = filepath.Join(pathSoFar, part)
		}
		if _, err := os.Stat(pathSoFar); os.IsNotExist(err) {
			//need to create it.
			os.Mkdir(pathSoFar, 0700)
		}
	}

	fileName := filepath.Join(pathSoFar, fileNamePart)

	file, err := os.Create(fileName)

	if err != nil {
		log.Println(err)
		return false
	}

	defer file.Close()

	puzzleText := grid.DataString()

	n, err := io.WriteString(file, puzzleText)

	if err != nil {
		log.Println(err)
		return false
	} else {
		if n < len(puzzleText) {
			log.Println("Didn't write full file, only wrote", n, "bytes of", len(puzzleText))
			return false
		}
	}
	return true
}

func vendPuzzle(min float64, max float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64) *sudoku.Grid {

	directory := filepath.Join(puzzleDirectoryParts(symmetryType, symmetryPercentage)...)

	if files, err := ioutil.ReadDir(directory); os.IsNotExist(err) {
		//The directory doesn't exist.
		return nil
	} else {
		//OK, the directory exists, now see which puzzles are there and if any fit. If one does, vend it and delete the file.
		for _, file := range files {
			//See what this actually returns.
			filenameParts := strings.Split(file.Name(), ".")

			//Remember: there's a dot in the filename due to the float seperator.
			//TODO: shouldn't "sdk" be in a constant somewhere?
			if len(filenameParts) != 3 || filenameParts[2] != "sdk" {
				continue
			}

			difficulty, err := strconv.ParseFloat(strings.Join(filenameParts[0:2], "."), 64)
			if err != nil {
				continue
			}

			if min <= difficulty && difficulty <= max {
				//Found a puzzle!
				grid := sudoku.NewGrid()
				fullFileName := filepath.Join(directory, file.Name())
				grid.LoadFromFile(fullFileName)
				os.Remove(fullFileName)
				return grid
			}
		}
	}
	return nil
}

func generatePuzzle(min float64, max float64, symmetryType sudoku.SymmetryType, symmetryPercentage float64, skipCache bool) *sudoku.Grid {
	var result *sudoku.Grid

	if !skipCache {
		result = vendPuzzle(min, max, symmetryType, symmetryPercentage)

		if result != nil {
			log.Println("Vending a puzzle from the cache.")
			return result
		}
	}

	//We'll have to generate one ourselves.
	count := 0
	for {
		//The first time we don't bother saying what number attemp it is, because if the first run is likely to generate a useable puzzle it's just noise.
		if count != 0 {
			log.Println("Attempt", count, "at generating puzzle.")
		}

		result = sudoku.GenerateGrid(symmetryType, symmetryPercentage)

		difficulty := result.Difficulty()

		if difficulty >= min && difficulty <= max {
			return result
		}

		log.Println("Rejecting grid of difficulty", difficulty)
		if storePuzzle(result, difficulty, symmetryType, symmetryPercentage) {
			log.Println("Stored the puzzle for future use.")
		}

		count++
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

// Version is inserted at build using --ldflags -X
// But we don't really build that way any longer.
// So include a fallback version number that's not useless.
var Version = "0.8.1"

const socketName = "/var/run/edgectl.socket"
const logfile = "/tmp/edgectl.log"
const apiVersion = 1

var displayVersion = fmt.Sprintf("v%s (api v%d)", Version, apiVersion)

const failedToConnect = "Unable to connect to the daemon (See \"edgectl help daemon\")"

var daemonHelp = `The Edge Control Daemon is a long-lived background component that manages
connections and network state.

Launch the Edge Control Daemon:
    sudo edgectl daemon

Examine the Daemon's log output in
    ` + logfile + `
to troubleshoot problems.
`

// edgectl is the full path to the Edge Control binary
var edgectl string

func main() {
	// Figure out our executable and save it
	if executable, err := os.Executable(); err != nil {
		fmt.Fprintf(os.Stderr, "Internal error: %v", err)
		os.Exit(1)
	} else {
		edgectl = executable
	}

	rootCmd := getRootCommand()
	err := rootCmd.Execute()
	if err != nil {
		os.Exit(1)
	}
}

func getRootCommand() *cobra.Command {
	myName := "Edge Control"
	if !isServerRunning() {
		myName = "Edge Control (daemon unavailable)"
	}

	rootCmd := &cobra.Command{
		Use:          "edgectl",
		Short:        myName,
		SilenceUsage: true, // https://github.com/spf13/cobra/issues/340
	}

	// Hidden/internal commands. These are called by Edge Control itself from
	// the correct context and execute in-place immediately.

	rootCmd.AddCommand(&cobra.Command{
		Use:    "daemon-foreground",
		Short:  "Launch Edge Control Daemon in the foreground (debug)",
		Args:   cobra.ExactArgs(0),
		Hidden: true,
		RunE: func(_ *cobra.Command, _ []string) error {
			return RunAsDaemon()
		},
	})
	teleproxyCmd := &cobra.Command{
		Use:    "teleproxy",
		Short:  "Impersonate Teleproxy (for internal use)",
		Hidden: true,
	}
	teleproxyCmd.AddCommand(&cobra.Command{
		Use:    "intercept",
		Short:  "Impersonate Teleproxy Intercept (for internal use)",
		Args:   cobra.NoArgs,
		Hidden: true,
		RunE: func(_ *cobra.Command, _ []string) error {
			return RunAsTeleproxyIntercept()
		},
	})
	teleproxyCmd.AddCommand(&cobra.Command{
		Use:    "bridge",
		Short:  "Impersonate Teleproxy Bridge (for internal use)",
		Args:   cobra.ExactArgs(2),
		Hidden: true,
		RunE: func(_ *cobra.Command, args []string) error {
			return RunAsTeleproxyBridge(args[0], args[1])
		},
	})
	rootCmd.AddCommand(teleproxyCmd)

	// Client commands. These are never sent to the daemon.

	rootCmd.AddCommand(&cobra.Command{
		Use:   "daemon",
		Short: "Launch Edge Control Daemon in the background (sudo)",
		Long:  daemonHelp,
		Args:  cobra.ExactArgs(0),
		RunE:  launchDaemon,
	})
	rootCmd.AddCommand(&cobra.Command{
		Use:   "login",
		Short: "Access the Ambassador Edge Stack admin UI",
		Args:  cobra.MaximumNArgs(1),
		RunE:  aesLogin,
	})
	licenseCmd := &cobra.Command{
		Use:   "license [flags] LICENSE_KEY",
		Short: "Set or update the Ambassador Edge Stack license key",
		Args:  cobra.ExactArgs(1),
		RunE:  aesLicense,
	}
	_ = licenseCmd.Flags().StringP(
		"context", "c", "",
		"The Kubernetes context to use. Defaults to the current kubectl context.",
	)
	_ = licenseCmd.Flags().StringP(
		"namespace", "n", "ambassador",
		"The Kubernetes namespace to use. Defaults to ambassador.",
	)
	rootCmd.AddCommand(licenseCmd)

	// Daemon commands. These should be forwarded to the daemon.

	nilDaemon := &Daemon{}
	daemonCmd := nilDaemon.getRootCommand(nil, nil, nil)
	walkSubcommands(daemonCmd)
	rootCmd.AddCommand(daemonCmd.Commands()...)
	rootCmd.PersistentFlags().AddFlagSet(daemonCmd.PersistentFlags())

	return rootCmd
}

func walkSubcommands(cmd *cobra.Command) {
	for _, subCmd := range cmd.Commands() {
		walkSubcommands(subCmd)
	}
	if cmd.RunE != nil {
		cmd.RunE = forwardToDaemon
	}
}

func forwardToDaemon(cmd *cobra.Command, _ []string) error {
	err := mainViaDaemon()
	if err != nil {
		// The version command is special because it must emit the client
		// version if the daemon is unavailable.
		if cmd.Use == "version" {
			fmt.Println("Client", displayVersion)
		}
		fmt.Println(failedToConnect)
	}
	return err
}

func launchDaemon(ccmd *cobra.Command, _ []string) error {
	if os.Geteuid() != 0 {
		fmt.Println("Edge Control Daemon must be launched as root.")
		fmt.Printf("\n  sudo %s\n\n", ccmd.CommandPath())
		return errors.New("root privileges required")
	}
	fmt.Println("Launching Edge Control Daemon", displayVersion)

	cmd := exec.Command(edgectl, "daemon-foreground")
	cmd.Env = os.Environ()
	cmd.Stdin = nil
	cmd.Stdout = nil
	cmd.Stderr = nil
	cmd.ExtraFiles = nil
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
	err := cmd.Start()
	if err != nil {
		return errors.Wrap(err, "failed to launch the server")
	}

	success := false
	for count := 0; count < 40; count++ {
		if isServerRunning() {
			success = true
			break
		}
		if count == 4 {
			fmt.Println("Waiting for daemon to start...")
		}
		time.Sleep(250 * time.Millisecond)
	}
	if !success {
		fmt.Println("Server did not come up!")
		fmt.Printf("Take a look at %s for more information.\n", logfile)
		return errors.New("launch failed")
	}
	return nil
}
<commit_msg>Revert "Give Edge Control a version number when built directly"<commit_after>package main

import (
	"fmt"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/pkg/errors"
	"github.com/spf13/cobra"
)

// Version is inserted at build using --ldflags -X
var Version = "(unknown version)"

const socketName = "/var/run/edgectl.socket"
const logfile = "/tmp/edgectl.log"
const apiVersion = 1

var displayVersion = fmt.Sprintf("v%s (api v%d)", Version, apiVersion)

const failedToConnect = "Unable to connect to the daemon (See \"edgectl help daemon\")"

var daemonHelp = `The Edge Control Daemon is a long-lived background component that manages
connections and network state.

Launch the Edge Control Daemon:
    sudo edgectl daemon

Examine the Daemon's log output in
    ` + logfile + `
to troubleshoot problems.
`

// edgectl is the full path to the Edge Control binary
var edgectl string

func main() {
	// Figure out our executable and save it
	if executable, err := os.Executable(); err != nil {
		fmt.Fprintf(os.Stderr, "Internal error: %v", err)
		os.Exit(1)
	} else {
		edgectl = executable
	}

	rootCmd := getRootCommand()
	err := rootCmd.Execute()
	if err != nil {
		os.Exit(1)
	}
}

func getRootCommand() *cobra.Command {
	myName := "Edge Control"
	if !isServerRunning() {
		myName = "Edge Control (daemon unavailable)"
	}

	rootCmd := &cobra.Command{
		Use:          "edgectl",
		Short:        myName,
		SilenceUsage: true, // https://github.com/spf13/cobra/issues/340
	}

	// Hidden/internal commands. These are called by Edge Control itself from
	// the correct context and execute in-place immediately.

	rootCmd.AddCommand(&cobra.Command{
		Use:    "daemon-foreground",
		Short:  "Launch Edge Control Daemon in the foreground (debug)",
		Args:   cobra.ExactArgs(0),
		Hidden: true,
		RunE: func(_ *cobra.Command, _ []string) error {
			return RunAsDaemon()
		},
	})
	teleproxyCmd := &cobra.Command{
		Use:    "teleproxy",
		Short:  "Impersonate Teleproxy (for internal use)",
		Hidden: true,
	}
	teleproxyCmd.AddCommand(&cobra.Command{
		Use:    "intercept",
		Short:  "Impersonate Teleproxy Intercept (for internal use)",
		Args:   cobra.NoArgs,
		Hidden: true,
		RunE: func(_ *cobra.Command, _ []string) error {
			return RunAsTeleproxyIntercept()
		},
	})
	teleproxyCmd.AddCommand(&cobra.Command{
		Use:    "bridge",
		Short:  "Impersonate Teleproxy Bridge (for internal use)",
		Args:   cobra.ExactArgs(2),
		Hidden: true,
		RunE: func(_ *cobra.Command, args []string) error {
			return RunAsTeleproxyBridge(args[0], args[1])
		},
	})
	rootCmd.AddCommand(teleproxyCmd)

	// Client commands. These are never sent to the daemon.

	rootCmd.AddCommand(&cobra.Command{
		Use:   "daemon",
		Short: "Launch Edge Control Daemon in the background (sudo)",
		Long:  daemonHelp,
		Args:  cobra.ExactArgs(0),
		RunE:  launchDaemon,
	})
	rootCmd.AddCommand(&cobra.Command{
		Use:   "login",
		Short: "Access the Ambassador Edge Stack admin UI",
		Args:  cobra.MaximumNArgs(1),
		RunE:  aesLogin,
	})
	licenseCmd := &cobra.Command{
		Use:   "license [flags] LICENSE_KEY",
		Short: "Set or update the Ambassador Edge Stack license key",
		Args:  cobra.ExactArgs(1),
		RunE:  aesLicense,
	}
	_ = licenseCmd.Flags().StringP(
		"context", "c", "",
		"The Kubernetes context to use. Defaults to the current kubectl context.",
	)
	_ = licenseCmd.Flags().StringP(
		"namespace", "n", "ambassador",
		"The Kubernetes namespace to use. Defaults to ambassador.",
	)
	rootCmd.AddCommand(licenseCmd)

	// Daemon commands. These should be forwarded to the daemon.

	nilDaemon := &Daemon{}
	daemonCmd := nilDaemon.getRootCommand(nil, nil, nil)
	walkSubcommands(daemonCmd)
	rootCmd.AddCommand(daemonCmd.Commands()...)
	rootCmd.PersistentFlags().AddFlagSet(daemonCmd.PersistentFlags())

	return rootCmd
}

func walkSubcommands(cmd *cobra.Command) {
	for _, subCmd := range cmd.Commands() {
		walkSubcommands(subCmd)
	}
	if cmd.RunE != nil {
		cmd.RunE = forwardToDaemon
	}
}

func forwardToDaemon(cmd *cobra.Command, _ []string) error {
	err := mainViaDaemon()
	if err != nil {
		// The version command is special because it must emit the client
		// version if the daemon is unavailable.
		if cmd.Use == "version" {
			fmt.Println("Client", displayVersion)
		}
		fmt.Println(failedToConnect)
	}
	return err
}

func launchDaemon(ccmd *cobra.Command, _ []string) error {
	if os.Geteuid() != 0 {
		fmt.Println("Edge Control Daemon must be launched as root.")
		fmt.Printf("\n  sudo %s\n\n", ccmd.CommandPath())
		return errors.New("root privileges required")
	}
	fmt.Println("Launching Edge Control Daemon", displayVersion)

	cmd := exec.Command(edgectl, "daemon-foreground")
	cmd.Env = os.Environ()
	cmd.Stdin = nil
	cmd.Stdout = nil
	cmd.Stderr = nil
	cmd.ExtraFiles = nil
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
	err := cmd.Start()
	if err != nil {
		return errors.Wrap(err, "failed to launch the server")
	}

	success := false
	for count := 0; count < 40; count++ {
		if isServerRunning() {
			success = true
			break
		}
		if count == 4 {
			fmt.Println("Waiting for daemon to start...")
		}
		time.Sleep(250 * time.Millisecond)
	}
	if !success {
		fmt.Println("Server did not come up!")
		fmt.Printf("Take a look at %s for more information.\n", logfile)
		return errors.New("launch failed")
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright The Helm 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 (
	"errors"
	"fmt"
	"io"

	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"

	apiVersion "k8s.io/apimachinery/pkg/version"
	"k8s.io/helm/pkg/helm"
	pb "k8s.io/helm/pkg/proto/hapi/version"
	"k8s.io/helm/pkg/version"
)

const versionDesc = `
Show the client and server versions for Helm and tiller.

This will print a representation of the client and server versions of Helm and
Tiller. The output will look something like this:

Client: &version.Version{SemVer:"v2.0.0", GitCommit:"ff52399e51bb880526e9cd0ed8386f6433b74da1", GitTreeState:"clean"}
Server: &version.Version{SemVer:"v2.0.0", GitCommit:"b0c113dfb9f612a9add796549da66c0d294508a3", GitTreeState:"clean"}

- SemVer is the semantic version of the release.
- GitCommit is the SHA for the commit that this version was built from.
- GitTreeState is "clean" if there are no local code changes when this binary was
  built, and "dirty" if the binary was built from locally modified code.

To print just the client version, use '--client'. To print just the server version,
use '--server'.
`

type versionCmd struct {
	out        io.Writer
	client     helm.Interface
	showClient bool
	showServer bool
	short      bool
	template   string
}

func newVersionCmd(c helm.Interface, out io.Writer) *cobra.Command {
	version := &versionCmd{
		client: c,
		out:    out,
	}

	cmd := &cobra.Command{
		Use:   "version",
		Short: "print the client/server version information",
		Long:  versionDesc,
		RunE: func(cmd *cobra.Command, args []string) error {
			// If neither is explicitly set, show both.
			if !version.showClient && !version.showServer {
				version.showClient, version.showServer = true, true
			}
			if version.showServer {
				// We do this manually instead of in PreRun because we only
				// need a tunnel if server version is requested.
				err := setupConnection()
				if err != nil {
					return err
				}
			}
			version.client = ensureHelmClient(version.client)
			return version.run()
		},
	}
	f := cmd.Flags()
	f.BoolVarP(&version.showClient, "client", "c", false, "client version only")
	f.BoolVarP(&version.showServer, "server", "s", false, "server version only")
	f.BoolVar(&version.short, "short", false, "print the version number")
	f.StringVar(&version.template, "template", "", "template for version string format")

	return cmd
}

func (v *versionCmd) run() error {
	// Store map data for template rendering
	data := map[string]interface{}{}

	if v.showClient {
		cv := version.GetVersionProto()
		if v.template != "" {
			data["Client"] = cv
		} else {
			fmt.Fprintf(v.out, "Client: %s\n", formatVersion(cv, v.short))
		}
	}

	if !v.showServer {
		return tpl(v.template, data, v.out)
	}

	if settings.Debug {
		k8sVersion, err := getK8sVersion()
		if err != nil {
			return err
		}
		fmt.Fprintf(v.out, "Kubernetes: %#v\n", k8sVersion)
	}
	resp, err := v.client.GetVersion()
	if err != nil {
		if grpc.Code(err) == codes.Unimplemented {
			return errors.New("server is too old to know its version")
		}
		debug("%s", err)
		return errors.New("cannot connect to Tiller")
	}

	if v.template != "" {
		data["Server"] = resp.Version
	} else {
		fmt.Fprintf(v.out, "Server: %s\n", formatVersion(resp.Version, v.short))
	}
	return tpl(v.template, data, v.out)
}

func getK8sVersion() (*apiVersion.Info, error) {
	var v *apiVersion.Info
	_, client, err := getKubeClient(settings.KubeContext, settings.KubeConfig)
	if err != nil {
		return v, err
	}
	v, err = client.Discovery().ServerVersion()
	return v, err
}

func formatVersion(v *pb.Version, short bool) string {
	if short {
		return fmt.Sprintf("%s+g%s", v.SemVer, v.GitCommit[:7])
	}
	return fmt.Sprintf("%#v", v)
}
<commit_msg>setup connection after displaying client version<commit_after>/*
Copyright The Helm 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 (
	"errors"
	"fmt"
	"io"

	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"

	apiVersion "k8s.io/apimachinery/pkg/version"
	"k8s.io/helm/pkg/helm"
	pb "k8s.io/helm/pkg/proto/hapi/version"
	"k8s.io/helm/pkg/version"
)

const versionDesc = `
Show the client and server versions for Helm and tiller.

This will print a representation of the client and server versions of Helm and
Tiller. The output will look something like this:

Client: &version.Version{SemVer:"v2.0.0", GitCommit:"ff52399e51bb880526e9cd0ed8386f6433b74da1", GitTreeState:"clean"}
Server: &version.Version{SemVer:"v2.0.0", GitCommit:"b0c113dfb9f612a9add796549da66c0d294508a3", GitTreeState:"clean"}

- SemVer is the semantic version of the release.
- GitCommit is the SHA for the commit that this version was built from.
- GitTreeState is "clean" if there are no local code changes when this binary was
  built, and "dirty" if the binary was built from locally modified code.

To print just the client version, use '--client'. To print just the server version,
use '--server'.
`

type versionCmd struct {
	out        io.Writer
	client     helm.Interface
	showClient bool
	showServer bool
	short      bool
	template   string
}

func newVersionCmd(c helm.Interface, out io.Writer) *cobra.Command {
	version := &versionCmd{
		client: c,
		out:    out,
	}

	cmd := &cobra.Command{
		Use:   "version",
		Short: "print the client/server version information",
		Long:  versionDesc,
		RunE: func(cmd *cobra.Command, args []string) error {
			// If neither is explicitly set, show both.
			if !version.showClient && !version.showServer {
				version.showClient, version.showServer = true, true
			}
			return version.run()
		},
	}
	f := cmd.Flags()
	f.BoolVarP(&version.showClient, "client", "c", false, "client version only")
	f.BoolVarP(&version.showServer, "server", "s", false, "server version only")
	f.BoolVar(&version.short, "short", false, "print the version number")
	f.StringVar(&version.template, "template", "", "template for version string format")

	return cmd
}

func (v *versionCmd) run() error {
	// Store map data for template rendering
	data := map[string]interface{}{}

	if v.showClient {
		cv := version.GetVersionProto()
		if v.template != "" {
			data["Client"] = cv
		} else {
			fmt.Fprintf(v.out, "Client: %s\n", formatVersion(cv, v.short))
		}
	}

	if !v.showServer {
		return tpl(v.template, data, v.out)
	}

	// We do this manually instead of in PreRun because we only
	// need a tunnel if server version is requested.
	if err := setupConnection(); err != nil {
		return err
	}
	v.client = ensureHelmClient(v.client)

	if settings.Debug {
		k8sVersion, err := getK8sVersion()
		if err != nil {
			return err
		}
		fmt.Fprintf(v.out, "Kubernetes: %#v\n", k8sVersion)
	}
	resp, err := v.client.GetVersion()
	if err != nil {
		if grpc.Code(err) == codes.Unimplemented {
			return errors.New("server is too old to know its version")
		}
		debug("%s", err)
		return errors.New("cannot connect to Tiller")
	}

	if v.template != "" {
		data["Server"] = resp.Version
	} else {
		fmt.Fprintf(v.out, "Server: %s\n", formatVersion(resp.Version, v.short))
	}
	return tpl(v.template, data, v.out)
}

func getK8sVersion() (*apiVersion.Info, error) {
	var v *apiVersion.Info
	_, client, err := getKubeClient(settings.KubeContext, settings.KubeConfig)
	if err != nil {
		return v, err
	}
	v, err = client.Discovery().ServerVersion()
	return v, err
}

func formatVersion(v *pb.Version, short bool) string {
	if short {
		return fmt.Sprintf("%s+g%s", v.SemVer, v.GitCommit[:7])
	}
	return fmt.Sprintf("%#v", v)
}
<|endoftext|>


<commit_before>// Copyright © 2016 Govinda Fichtner <govinda.fichtner@googlemail.com>
//
// 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.

package cmd

import (
	"fmt"
	"github.com/spf13/cobra"
	"io/ioutil"
	"os/exec"
	"regexp"
	"strings"
)

// setCmd represents the set command
var setHostnameCmd = &cobra.Command{
	Use:   "set [hostname]",
	Short: "Set a hostname",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) > 0 {
			set_hostname(args[0])
		} else {
			cmd.Help()
		}
	},
}

func init() {
	hostnameCmd.AddCommand(setHostnameCmd)
}

func set_hostname(args ...string) {
	hostname := ""

	// if we have hostname in config file use that
	if config.IsSet("hostname") {
		hostname = config.GetString("hostname")
	}

	// if we have a hostname as command line arg use that
	if len(args) > 0 {
		hostname = args[0]
	}

	if hostname == "" && cfgFile == "" {
		fmt.Println("missing hostname argument")
		return
	}

	if hostname != "" {
		err := ioutil.WriteFile("/etc/hostname", []byte(hostname), 0644)
		if err != nil {
			panic(err)
		}

		input, err := ioutil.ReadFile("/etc/hosts")
		if err != nil {
			panic(err)
		}

		lines := strings.Split(string(input), "\n")

		for i, line := range lines {
			if strings.Contains(line, "127.0.0.1	localhost") {
				lines[i] = fmt.Sprintf("127.0.0.1	localhost	%s", hostname)
			}
		}
		output := strings.Join(lines, "\n")
		err = ioutil.WriteFile("/etc/hosts", []byte(output), 0644)
		if err != nil {
			panic(err)
		}

		err = exec.Command("hostname", hostname).Run()
		if err != nil {
			fmt.Println("Unable to set hostname: ", err)
		}

		// ensure that dhcp server and avahi daemon are aware of new hostname
		for _, interfaceName := range activeInterfaces() {
			err = exec.Command("/sbin/ifdown", interfaceName).Run()
			if err != nil {
				fmt.Println("Unable to bring interface down: ", interfaceName, err)
			}

			err = exec.Command("/sbin/ifup", interfaceName).Run()
			if err != nil {
				fmt.Println("Unable to bring interface up: ", interfaceName, err)
			}
		}

		err = exec.Command("/bin/systemctl", "restart", "avahi-daemon").Run()
		if err != nil {
			fmt.Println("Unable to set hostname: ", err)
		}

		fmt.Printf("Set hostname: %s\n", hostname)
	}
}

func activeInterfaces() []string {
	var result []string
	output, err := exec.Command("ip", "link").Output()
	if err != nil {
		fmt.Println("Could not run 'ip link'", err)
	}
	lines := strings.Split(string(output), "\n")
	for _, line := range lines {
		interfaceIsUp, _ := regexp.MatchString("state UP", line)
		if interfaceIsUp {
			re := regexp.MustCompile(`^\d*:\s([a-z0-9@]*):`)
			result = append(result, re.FindStringSubmatch(line)[1])
		}
	}
	return result
}
<commit_msg>set hostname on its own line in /etc/hosts<commit_after>// Copyright © 2016 Govinda Fichtner <govinda.fichtner@googlemail.com>
//
// 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.

package cmd

import (
	"fmt"
	"github.com/spf13/cobra"
	"io/ioutil"
	"os/exec"
	"regexp"
	"strings"
)

// setCmd represents the set command
var setHostnameCmd = &cobra.Command{
	Use:   "set [hostname]",
	Short: "Set a hostname",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		if len(args) > 0 {
			set_hostname(args[0])
		} else {
			cmd.Help()
		}
	},
}

func init() {
	hostnameCmd.AddCommand(setHostnameCmd)
}

func set_hostname(args ...string) {
	hostname := ""

	// if we have hostname in config file use that
	if config.IsSet("hostname") {
		hostname = config.GetString("hostname")
	}

	// if we have a hostname as command line arg use that
	if len(args) > 0 {
		hostname = args[0]
	}

	if hostname == "" && cfgFile == "" {
		fmt.Println("missing hostname argument")
		return
	}

	if hostname != "" {
		err := ioutil.WriteFile("/etc/hostname", []byte(hostname), 0644)
		if err != nil {
			panic(err)
		}

		input, err := ioutil.ReadFile("/etc/hosts")
		if err != nil {
			panic(err)
		}

		lines := strings.Split(string(input), "\n")
		lines_new := []string{}

		for i, line := range lines {
			if strings.Contains(line, "127.0.0.1	localhost") {
				lines_new = append(lines_new, lines[0:i+1]...)
				lines_new = append(lines_new, fmt.Sprintf("127.0.0.1	%s", hostname))
				lines_new = append(lines_new, lines[i+1:]...)
			}
		}
		output := strings.Join(lines_new, "\n")
		err = ioutil.WriteFile("/etc/hosts", []byte(output), 0644)
		if err != nil {
			panic(err)
		}

		err = exec.Command("hostname", hostname).Run()
		if err != nil {
			fmt.Println("Unable to set hostname: ", err)
		}

		// ensure that dhcp server and avahi daemon are aware of new hostname
		for _, interfaceName := range activeInterfaces() {
			err = exec.Command("/sbin/ifdown", interfaceName).Run()
			if err != nil {
				fmt.Println("Unable to bring interface down: ", interfaceName, err)
			}

			err = exec.Command("/sbin/ifup", interfaceName).Run()
			if err != nil {
				fmt.Println("Unable to bring interface up: ", interfaceName, err)
			}
		}

		err = exec.Command("/bin/systemctl", "restart", "avahi-daemon").Run()
		if err != nil {
			fmt.Println("Unable to set hostname: ", err)
		}

		fmt.Printf("Set hostname: %s\n", hostname)
	}
}

func activeInterfaces() []string {
	var result []string
	output, err := exec.Command("ip", "link").Output()
	if err != nil {
		fmt.Println("Could not run 'ip link'", err)
	}
	lines := strings.Split(string(output), "\n")
	for _, line := range lines {
		interfaceIsUp, _ := regexp.MatchString("state UP", line)
		if interfaceIsUp {
			re := regexp.MustCompile(`^\d*:\s([a-z0-9@]*):`)
			result = append(result, re.FindStringSubmatch(line)[1])
		}
	}
	return result
}
<|endoftext|>


<commit_before>//
// Copyright 2021 The Ent 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"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"sync"

	"cloud.google.com/go/firestore"
	"github.com/google/ent/index"
	"github.com/google/ent/utils"
	"github.com/spf13/cobra"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

var (
	indexFlag string

	firebaseCredentials string
	firebaseProject     string
	concurrency         int

	urlFlag string
)

type URL struct {
	URL string
}

func server(cmd *cobra.Command, args []string) {
	ctx := context.Background()
	client, err := firestore.NewClient(ctx, firebaseProject, option.WithCredentialsFile(firebaseCredentials))
	if err != nil {
		log.Fatalf("could not create client: %v", err)
	}
	if indexFlag == "" {
		log.Fatal("index flag is required")
	}
	iter := client.Collection("urls").Documents(ctx)
	defer iter.Stop()
	wg := sync.WaitGroup{}
	tokens := make(chan struct{}, concurrency)
	for {
		doc, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			log.Fatalf("error iterating over URLs: %v", err)
		}
		url := doc.Data()["url"].(string)
		if url == "" {
			continue
		}
		wg.Add(1)
		tokens <- struct{}{}
		go func() {
			defer func() {
				wg.Done()
				<-tokens
			}()
			err := fetch(url)
			if err != nil {
				log.Printf("could not fetch URL: %v", err)
			}
		}()
	}
	wg.Wait()
}

func fetchCmd(cmd *cobra.Command, args []string) {
	err := fetch(urlFlag)
	if err != nil {
		log.Fatalf("could not fetch URL: %v", err)
	}
}

func fetch(urlString string) error {
	log.Print(urlString)
	parsedURL, err := url.Parse(urlString)
	if err != nil {
		return fmt.Errorf("invalid URL: %w", err)
	}

	if parsedURL.User != nil {
		return fmt.Errorf("non-empty user in URL")
	}
	if parsedURL.Fragment != "" {
		return fmt.Errorf("non-empty fragment in URL")
	}

	urlString = parsedURL.String()
	log.Printf("fetching %q", urlString)

	res, err := http.Get(urlString)
	if err != nil {
		return fmt.Errorf("could not fetch URL %q: %w", urlString, err)
	}
	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("invalid error code %d (%s)", res.StatusCode, res.Status)
	}

	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return fmt.Errorf("could not read HTTP body: %w", err)
	}
	h := utils.ComputeHash(data)

	l := filepath.Join(indexFlag, index.HashToPath(h), index.EntryFilename)

	var e index.IndexEntry
	if _, err := os.Stat(l); err == nil {
		log.Printf("index entry existing: %q", l)
		bytes, err := ioutil.ReadFile(l)
		if err != nil {
			return fmt.Errorf("could not read index entry: %w", err)
		}
		err = json.Unmarshal(bytes, &e)
		if err != nil {
			return fmt.Errorf("could not unmarshal JSON for index entry: %w", err)
		}
		if sort.SearchStrings(e.URLS, urlString) != -1 {
			log.Printf("URL already indexed in entry: %+v", e)
			// Nothing to do.
			return nil
		}
		e.URLS = append(e.URLS, urlString)
		sort.Strings(e.URLS)
	} else {
		e = index.IndexEntry{
			Hash: h,
			Size: len(data),
			URLS: []string{urlString},
		}
	}
	log.Printf("index entry to create: %+v", e)

	es, err := json.Marshal(e)
	if err != nil {
		return fmt.Errorf("could not marshal JSON: %w", err)
	}
	err = os.MkdirAll(filepath.Dir(l), 0755)
	if err != nil {
		return fmt.Errorf("could not create file: %w", err)
	}
	err = ioutil.WriteFile(l, es, 0644)
	if err != nil {
		return fmt.Errorf("could not write to file: %w", err)
	}

	log.Printf("index entry created: %q", l)

	return nil
}

func main() {
	var rootCmd = &cobra.Command{Use: "indexer"}

	serverCmd :=
		&cobra.Command{
			Use:   "server",
			Short: "Run the indexer server",
			Run:   server,
		}
	serverCmd.PersistentFlags().StringVar(&indexFlag, "index", "", "path to index repository")
	serverCmd.PersistentFlags().StringVar(&firebaseProject, "firebase-project", "", "Firebase project name")
	serverCmd.PersistentFlags().StringVar(&firebaseCredentials, "firebase-credentials", "", "file with Firebase credentials")
	serverCmd.PersistentFlags().IntVar(&concurrency, "concurrency", 10, "HTTP fetch concurrency")
	rootCmd.AddCommand(serverCmd)

	getCmd := &cobra.Command{
		Use:   "fetch",
		Short: "Fetch and index a single URL",
		Run:   fetchCmd,
	}
	getCmd.PersistentFlags().StringVar(&indexFlag, "index", "", "path to index repository")
	getCmd.PersistentFlags().StringVar(&urlFlag, "url", "", "url of the entry to index")
	rootCmd.AddCommand(getCmd)

	rootCmd.Execute()
}
<commit_msg>Better formatting<commit_after>//
// Copyright 2021 The Ent 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"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"sort"
	"sync"

	"cloud.google.com/go/firestore"
	"github.com/google/ent/index"
	"github.com/google/ent/utils"
	"github.com/spf13/cobra"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

var (
	indexFlag string

	firebaseCredentials string
	firebaseProject     string
	concurrency         int

	urlFlag string
)

type URL struct {
	URL string
}

func server(cmd *cobra.Command, args []string) {
	ctx := context.Background()
	client, err := firestore.NewClient(ctx, firebaseProject, option.WithCredentialsFile(firebaseCredentials))
	if err != nil {
		log.Fatalf("could not create client: %v", err)
	}
	if indexFlag == "" {
		log.Fatal("index flag is required")
	}
	iter := client.Collection("urls").Documents(ctx)
	defer iter.Stop()
	wg := sync.WaitGroup{}
	tokens := make(chan struct{}, concurrency)
	for {
		doc, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			log.Fatalf("error iterating over URLs: %v", err)
		}
		url := doc.Data()["url"].(string)
		if url == "" {
			continue
		}
		wg.Add(1)
		tokens <- struct{}{}
		go func() {
			defer func() {
				wg.Done()
				<-tokens
			}()
			_, err := fetch(url)
			if err != nil {
				log.Printf("could not fetch URL: %v", err)
			}
		}()
	}
	wg.Wait()
}

func fetchCmd(cmd *cobra.Command, args []string) {
	e, err := fetch(urlFlag)
	if err != nil {
		log.Fatalf("could not fetch URL: %v", err)
	}
	// Print hash to stdout.
	fmt.Printf("%s\n", e.Hash)
}

func fetch(urlString string) (*index.IndexEntry, error) {
	log.Print(urlString)
	parsedURL, err := url.Parse(urlString)
	if err != nil {
		return nil, fmt.Errorf("invalid URL: %w", err)
	}

	if parsedURL.User != nil {
		return nil, fmt.Errorf("non-empty user in URL")
	}
	if parsedURL.Fragment != "" {
		return nil, fmt.Errorf("non-empty fragment in URL")
	}

	urlString = parsedURL.String()
	log.Printf("fetching %q", urlString)

	res, err := http.Get(urlString)
	if err != nil {
		return nil, fmt.Errorf("could not fetch URL %q: %w", urlString, err)
	}
	if res.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("invalid error code %d (%s)", res.StatusCode, res.Status)
	}

	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return nil, fmt.Errorf("could not read HTTP body: %w", err)
	}
	h := utils.ComputeHash(data)

	l := filepath.Join(indexFlag, index.HashToPath(h), index.EntryFilename)

	var e index.IndexEntry
	if _, err := os.Stat(l); err == nil {
		log.Printf("index entry existing: %q", l)
		bytes, err := ioutil.ReadFile(l)
		if err != nil {
			return nil, fmt.Errorf("could not read index entry: %w", err)
		}
		err = json.Unmarshal(bytes, &e)
		if err != nil {
			return nil, fmt.Errorf("could not unmarshal JSON for index entry: %w", err)
		}
		if sort.SearchStrings(e.URLS, urlString) != -1 {
			log.Printf("URL already indexed in entry: %+v", e)
			// Nothing to do.
			return &e, nil
		}
		e.URLS = append(e.URLS, urlString)
		sort.Strings(e.URLS)
	} else {
		e = index.IndexEntry{
			Hash: h,
			Size: len(data),
			URLS: []string{urlString},
		}
	}
	log.Printf("index entry to create: %+v", e)

	es, err := json.Marshal(e)
	if err != nil {
		return nil, fmt.Errorf("could not marshal JSON: %w", err)
	}
	err = os.MkdirAll(filepath.Dir(l), 0755)
	if err != nil {
		return nil, fmt.Errorf("could not create file: %w", err)
	}
	err = ioutil.WriteFile(l, es, 0644)
	if err != nil {
		return nil, fmt.Errorf("could not write to file: %w", err)
	}

	log.Printf("index entry created: %q", l)

	return &e, nil
}

func main() {
	var rootCmd = &cobra.Command{Use: "indexer"}

	serverCmd :=
		&cobra.Command{
			Use:   "server",
			Short: "Run the indexer server",
			Run:   server,
		}
	serverCmd.PersistentFlags().StringVar(&indexFlag, "index", "", "path to index repository")
	serverCmd.PersistentFlags().StringVar(&firebaseProject, "firebase-project", "", "Firebase project name")
	serverCmd.PersistentFlags().StringVar(&firebaseCredentials, "firebase-credentials", "", "file with Firebase credentials")
	serverCmd.PersistentFlags().IntVar(&concurrency, "concurrency", 10, "HTTP fetch concurrency")
	rootCmd.AddCommand(serverCmd)

	getCmd := &cobra.Command{
		Use:   "fetch",
		Short: "Fetch and index a single URL",
		Run:   fetchCmd,
	}
	getCmd.PersistentFlags().StringVar(&indexFlag, "index", "", "path to index repository")
	getCmd.PersistentFlags().StringVar(&urlFlag, "url", "", "url of the entry to index")
	rootCmd.AddCommand(getCmd)

	rootCmd.Execute()
}
<|endoftext|>