<commit_before>package main

import (
	"path/filepath"

	"github.com/pkg/errors"

	"github.com/havoc-io/mutagen/cmd"
	"github.com/havoc-io/mutagen/daemon"
	"github.com/havoc-io/mutagen/rpc"
	"github.com/havoc-io/mutagen/session"
	"github.com/havoc-io/mutagen/url"
)

var createUsage = `usage: mutagen create [-h|--help] [-i|--ignore=<pattern>]
                      <alpha> <beta>

Creates and starts a new synchronization session.
`

type ignorePatterns []string

func (p *ignorePatterns) String() string {
	return "ignore patterns"
}

func (p *ignorePatterns) Set(value string) error {
	*p = append(*p, value)
	return nil
}

func createMain(arguments []string) error {
	// Parse command line arguments.
	var ignores ignorePatterns
	flagSet := cmd.NewFlagSet("create", createUsage, []int{2})
	flagSet.VarP(&ignores, "ignore", "i", "specify ignore paths")
	urls := flagSet.ParseOrDie(arguments)

	// Extract and parse URLs.
	alpha, err := url.Parse(urls[0])
	if err != nil {
		return errors.Wrap(err, "unable to parse alpha URL")
	}
	beta, err := url.Parse(urls[1])
	if err != nil {
		return errors.Wrap(err, "unable to parse beta URL")
	}

	// If either URL is a relative path, convert it to an absolute path.
	if alpha.Protocol == url.Protocol_Local {
		if alphaPath, err := filepath.Abs(alpha.Path); err != nil {
			return errors.Wrap(err, "unable to make alpha path absolute")
		} else {
			alpha.Path = alphaPath
		}
	}
	if beta.Protocol == url.Protocol_Local {
		if betaPath, err := filepath.Abs(beta.Path); err != nil {
			return errors.Wrap(err, "unable to make beta path absolute")
		} else {
			beta.Path = betaPath
		}
	}

	// Create a daemon client.
	daemonClient := rpc.NewClient(daemon.NewOpener())

	// Invoke the session creation method and ensure the resulting stream is
	// closed when we're done.
	stream, err := daemonClient.Invoke(session.MethodCreate)
	if err != nil {
		return errors.Wrap(err, "unable to invoke session creation")
	}
	defer stream.Close()

	// Send the initial request.
	if err := stream.Send(session.CreateRequest{
		Alpha:   alpha,
		Beta:    beta,
		Ignores: []string(ignores),
	}); err != nil {
		return errors.Wrap(err, "unable to send creation request")
	}

	// Handle authentication challenges.
	return handleChallengePrompts(stream)
}
<commit_msg>Switched create command path normalization to be more consistent.<commit_after>package main

import (
	"github.com/pkg/errors"

	"github.com/havoc-io/mutagen/cmd"
	"github.com/havoc-io/mutagen/daemon"
	"github.com/havoc-io/mutagen/filesystem"
	"github.com/havoc-io/mutagen/rpc"
	"github.com/havoc-io/mutagen/session"
	"github.com/havoc-io/mutagen/url"
)

var createUsage = `usage: mutagen create [-h|--help] [-i|--ignore=<pattern>]
                      <alpha> <beta>

Creates and starts a new synchronization session.
`

type ignorePatterns []string

func (p *ignorePatterns) String() string {
	return "ignore patterns"
}

func (p *ignorePatterns) Set(value string) error {
	*p = append(*p, value)
	return nil
}

func createMain(arguments []string) error {
	// Parse command line arguments.
	var ignores ignorePatterns
	flagSet := cmd.NewFlagSet("create", createUsage, []int{2})
	flagSet.VarP(&ignores, "ignore", "i", "specify ignore paths")
	urls := flagSet.ParseOrDie(arguments)

	// Extract and parse URLs.
	alpha, err := url.Parse(urls[0])
	if err != nil {
		return errors.Wrap(err, "unable to parse alpha URL")
	}
	beta, err := url.Parse(urls[1])
	if err != nil {
		return errors.Wrap(err, "unable to parse beta URL")
	}

	// If either URL is a local path, make sure it's normalized.
	if alpha.Protocol == url.Protocol_Local {
		if alphaPath, err := filesystem.Normalize(alpha.Path); err != nil {
			return errors.Wrap(err, "unable to normalize alpha path")
		} else {
			alpha.Path = alphaPath
		}
	}
	if beta.Protocol == url.Protocol_Local {
		if betaPath, err := filesystem.Normalize(beta.Path); err != nil {
			return errors.Wrap(err, "unable to normalize beta path")
		} else {
			beta.Path = betaPath
		}
	}

	// Create a daemon client.
	daemonClient := rpc.NewClient(daemon.NewOpener())

	// Invoke the session creation method and ensure the resulting stream is
	// closed when we're done.
	stream, err := daemonClient.Invoke(session.MethodCreate)
	if err != nil {
		return errors.Wrap(err, "unable to invoke session creation")
	}
	defer stream.Close()

	// Send the initial request.
	if err := stream.Send(session.CreateRequest{
		Alpha:   alpha,
		Beta:    beta,
		Ignores: []string(ignores),
	}); err != nil {
		return errors.Wrap(err, "unable to send creation request")
	}

	// Handle authentication challenges.
	return handleChallengePrompts(stream)
}
<|endoftext|>


<commit_before>//+build darwin

package cmd

import (
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/rclone/rclone/fs/accounting"
)

// SigInfoHandler creates SigInfo handler
func SigInfoHandler() {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINFO)
	go func() {
		for range signals {
			log.Printf("%v\n", accounting.Stats)
		}
	}()
}
<commit_msg>cmd: fix up stats printing on macOS after accounting change<commit_after>//+build darwin

package cmd

import (
	"log"
	"os"
	"os/signal"
	"syscall"

	"github.com/rclone/rclone/fs/accounting"
)

// SigInfoHandler creates SigInfo handler
func SigInfoHandler() {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINFO)
	go func() {
		for range signals {
			log.Printf("%v\n", accounting.GlobalStats())
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/loadimpact/speedboat"
	"github.com/loadimpact/speedboat/simple"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"time"
)

// Configure the global logger.
func configureLogging(c *cli.Context) {
	log.SetLevel(log.InfoLevel)
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}
}

func parse(cc *cli.Context) (conf Config, err error) {
	switch len(cc.Args()) {
	case 0:
		if !cc.IsSet("script") && !cc.IsSet("url") {
			return conf, errors.New("No config file, script or URL")
		}
	case 1:
		bytes, err := ioutil.ReadFile(cc.Args()[0])
		if err != nil {
			return conf, errors.New("Couldn't read config file")
		}
		if err := yaml.Unmarshal(bytes, &conf); err != nil {
			return conf, errors.New("Couldn't parse config file")
		}
	default:
		return conf, errors.New("Too many arguments!")
	}

	// Let commandline flags override config files
	if cc.IsSet("script") {
		conf.Script = cc.String("script")
	}
	if cc.IsSet("url") {
		conf.URL = cc.String("url")
	}
	if cc.IsSet("vus") {
		conf.VUs = cc.Int("vus")
	}
	if cc.IsSet("duration") {
		conf.Duration = cc.Duration("duration").String()
	}

	return conf, nil
}

func dumpTest(t *speedboat.Test) {
	log.WithFields(log.Fields{
		"script": t.Script,
		"url":    t.URL,
	}).Info("General")
	for i, stage := range t.Stages {
		log.WithFields(log.Fields{
			"#":        i,
			"duration": stage.Duration,
			"start":    stage.StartVUs,
			"end":      stage.EndVUs,
		}).Info("Stage")
	}
}

func action(cc *cli.Context) error {
	conf, err := parse(cc)
	if err != nil {
		log.WithError(err).Fatal("Invalid arguments; see --help")
	}

	t, err := conf.MakeTest()
	if err != nil {
		log.WithError(err).Fatal("Configuration error")
	}

	if cc.Bool("dump") {
		dumpTest(&t)
		return nil
	}

	var runner speedboat.Runner
	switch {
	case t.URL != "":
		runner = simple.New()
	default:
		log.Fatal("No suitable runner found!")
	}

	ctx, _ := context.WithTimeout(context.Background(), t.TotalDuration())
	offset := time.Duration(0)
	for _, stage := range t.Stages {
		localOffset := offset
		go func() {
			time.Sleep(localOffset)
			c, _ := context.WithTimeout(ctx, stage.Duration)
			runner.RunVU(c, t)
		}()
		offset += stage.Duration
	}

	<-ctx.Done()

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"

	// Bootstrap using action-registered commandline flags
	app := cli.NewApp()
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "0.0.1a1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, v",
			Usage: "More verbose output",
		},
		cli.StringFlag{
			Name:  "script, s",
			Usage: "Script to run",
		},
		cli.StringFlag{
			Name:  "url",
			Usage: "URL to test",
		},
		cli.IntFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: 10,
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "Test duration",
			Value: time.Duration(10) * time.Second,
		},
		cli.StringFlag{
			Name:  "out-file, o",
			Usage: "Output raw metrics to a file",
		},
		cli.BoolFlag{
			Name:  "dump",
			Usage: "Dump parsed test and exit",
		},
	}
	app.Before = func(c *cli.Context) error {
		configureLogging(c)
		return nil
	}
	app.Action = action
	app.Run(os.Args)
}
<commit_msg>[docs] Comments explaining VU scheduling<commit_after>package main

import (
	"errors"
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/loadimpact/speedboat"
	"github.com/loadimpact/speedboat/simple"
	"golang.org/x/net/context"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"os"
	"time"
)

// Configure the global logger.
func configureLogging(c *cli.Context) {
	log.SetLevel(log.InfoLevel)
	if c.GlobalBool("verbose") {
		log.SetLevel(log.DebugLevel)
	}
}

func parse(cc *cli.Context) (conf Config, err error) {
	switch len(cc.Args()) {
	case 0:
		if !cc.IsSet("script") && !cc.IsSet("url") {
			return conf, errors.New("No config file, script or URL")
		}
	case 1:
		bytes, err := ioutil.ReadFile(cc.Args()[0])
		if err != nil {
			return conf, errors.New("Couldn't read config file")
		}
		if err := yaml.Unmarshal(bytes, &conf); err != nil {
			return conf, errors.New("Couldn't parse config file")
		}
	default:
		return conf, errors.New("Too many arguments!")
	}

	// Let commandline flags override config files
	if cc.IsSet("script") {
		conf.Script = cc.String("script")
	}
	if cc.IsSet("url") {
		conf.URL = cc.String("url")
	}
	if cc.IsSet("vus") {
		conf.VUs = cc.Int("vus")
	}
	if cc.IsSet("duration") {
		conf.Duration = cc.Duration("duration").String()
	}

	return conf, nil
}

func dumpTest(t *speedboat.Test) {
	log.WithFields(log.Fields{
		"script": t.Script,
		"url":    t.URL,
	}).Info("General")
	for i, stage := range t.Stages {
		log.WithFields(log.Fields{
			"#":        i,
			"duration": stage.Duration,
			"start":    stage.StartVUs,
			"end":      stage.EndVUs,
		}).Info("Stage")
	}
}

func action(cc *cli.Context) error {
	conf, err := parse(cc)
	if err != nil {
		log.WithError(err).Fatal("Invalid arguments; see --help")
	}

	t, err := conf.MakeTest()
	if err != nil {
		log.WithError(err).Fatal("Configuration error")
	}

	if cc.Bool("dump") {
		dumpTest(&t)
		return nil
	}

	// Inspect the test to find a suitable runner; additional ones can easily be added
	var runner speedboat.Runner
	switch {
	case t.URL != "":
		runner = simple.New()
	default:
		log.Fatal("No suitable runner found!")
	}

	// Schedule all configured VUs. Because we know the VU curves ahead of time, we:
	// - Make a context with the test's duration as timeout
	// - Loop through all the stages of the test
	// - Spawn VUs that:
	//     - Sleep until they're scheduled to start
	//     - Expire at the projected end of their lifecycles
	ctx, _ := context.WithTimeout(context.Background(), t.TotalDuration())
	offset := time.Duration(0)
	for _, stage := range t.Stages {
		localOffset := offset
		go func() {
			time.Sleep(localOffset)
			c, _ := context.WithTimeout(ctx, stage.Duration)
			runner.RunVU(c, t)
		}()
		offset += stage.Duration
	}

	// Wait until the end of the test
	<-ctx.Done()

	return nil
}

func main() {
	// Free up -v and -h for our own flags
	cli.VersionFlag.Name = "version"
	cli.HelpFlag.Name = "help, ?"

	// Bootstrap using action-registered commandline flags
	app := cli.NewApp()
	app.Name = "speedboat"
	app.Usage = "A next-generation load generator"
	app.Version = "0.0.1a1"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose, v",
			Usage: "More verbose output",
		},
		cli.StringFlag{
			Name:  "script, s",
			Usage: "Script to run",
		},
		cli.StringFlag{
			Name:  "url",
			Usage: "URL to test",
		},
		cli.IntFlag{
			Name:  "vus, u",
			Usage: "Number of VUs to simulate",
			Value: 10,
		},
		cli.DurationFlag{
			Name:  "duration, d",
			Usage: "Test duration",
			Value: time.Duration(10) * time.Second,
		},
		cli.StringFlag{
			Name:  "out-file, o",
			Usage: "Output raw metrics to a file",
		},
		cli.BoolFlag{
			Name:  "dump",
			Usage: "Dump parsed test and exit",
		},
	}
	app.Before = func(c *cli.Context) error {
		configureLogging(c)
		return nil
	}
	app.Action = action
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io"
	"io/ioutil"
	"log"
	"os/exec"
	"path/filepath"

	"github.com/mdlayher/tapcam/camera"
	"github.com/pkg/sftp"
)

var (
	device = flag.String("d", camera.DefaultDevice, "webcam device location")
	format = flag.String("f", string(camera.FormatJPEG), "webcam image capture format")
	size   = flag.String("s", camera.Resolution720p.String(), "webcam image size")
	host   = flag.String("host", "", "tapcamd host")
)

func main() {
	flag.Parse()

	resolution, err := camera.NewResolution(*size)
	if err != nil {
		log.Fatal(err)
	}

	cam, err := camera.New(
		camera.SetDevice(*device),
		camera.SetFormat(camera.Format(*format)),
		camera.SetResolution(resolution),
	)
	if err != nil {
		log.Fatal(err)
	}

	rc, done, err := cam.Capture()
	if err != nil {
		log.Fatal(err)
	}

	c, cdone, err := sftpClient(*host, ioutil.Discard)
	if err != nil {
		log.Fatal(err)
	}

	const (
		dir = "/tmp/tapcam"

		tmpName  = "latest.tmp.jpg"
		permName = "latest.jpg"
	)

	tmpFullName := filepath.Join(dir, tmpName)
	permFullName := filepath.Join(dir, permName)

	f, err := c.Create(tmpFullName)
	if err != nil {
		log.Fatal(err)
	}

	if _, err := io.Copy(f, rc); err != nil {
		log.Fatal(err)
	}

	if err := done(); err != nil {
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	if err := c.Remove(permFullName); err != nil {
		log.Fatal(err)
	}

	if err := c.Rename(tmpFullName, permFullName); err != nil {
		log.Fatal(err)
	}

	if err := c.Close(); err != nil {
		log.Fatal(err)
	}

	if err := cdone(); err != nil {
		log.Fatal(err)
	}
}

func sftpClient(host string, out io.Writer) (*sftp.Client, func() error, error) {
	// Connect to a remote host and request the sftp subsystem via the 'ssh'
	// command.  This assumes that passwordless login is correctly configured.
	cmd := exec.Command("ssh", host, "-s", "sftp")
	cmd.Stderr = out

	wr, err := cmd.StdinPipe()
	if err != nil {
		return nil, nil, err
	}
	rd, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}

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

	c, err := sftp.NewClientPipe(rd, wr)
	return c, func() error {
		return cmd.Wait()
	}, err
}
<commit_msg>tapcamctl: default to 1080p resolution<commit_after>package main

import (
	"flag"
	"io"
	"io/ioutil"
	"log"
	"os/exec"
	"path/filepath"

	"github.com/mdlayher/tapcam/camera"
	"github.com/pkg/sftp"
)

var (
	device = flag.String("d", camera.DefaultDevice, "webcam device location")
	format = flag.String("f", string(camera.FormatJPEG), "webcam image capture format")
	size   = flag.String("s", camera.Resolution1080p.String(), "webcam image size")
	host   = flag.String("host", "", "tapcamd host")
)

func main() {
	flag.Parse()

	resolution, err := camera.NewResolution(*size)
	if err != nil {
		log.Fatal(err)
	}

	cam, err := camera.New(
		camera.SetDevice(*device),
		camera.SetFormat(camera.Format(*format)),
		camera.SetResolution(resolution),
	)
	if err != nil {
		log.Fatal(err)
	}

	rc, done, err := cam.Capture()
	if err != nil {
		log.Fatal(err)
	}

	c, cdone, err := sftpClient(*host, ioutil.Discard)
	if err != nil {
		log.Fatal(err)
	}

	const (
		dir = "/tmp/tapcam"

		tmpName  = "latest.tmp.jpg"
		permName = "latest.jpg"
	)

	tmpFullName := filepath.Join(dir, tmpName)
	permFullName := filepath.Join(dir, permName)

	f, err := c.Create(tmpFullName)
	if err != nil {
		log.Fatal(err)
	}

	if _, err := io.Copy(f, rc); err != nil {
		log.Fatal(err)
	}

	if err := done(); err != nil {
		log.Fatal(err)
	}

	if err := f.Close(); err != nil {
		log.Fatal(err)
	}

	if err := c.Remove(permFullName); err != nil {
		log.Fatal(err)
	}

	if err := c.Rename(tmpFullName, permFullName); err != nil {
		log.Fatal(err)
	}

	if err := c.Close(); err != nil {
		log.Fatal(err)
	}

	if err := cdone(); err != nil {
		log.Fatal(err)
	}
}

func sftpClient(host string, out io.Writer) (*sftp.Client, func() error, error) {
	// Connect to a remote host and request the sftp subsystem via the 'ssh'
	// command.  This assumes that passwordless login is correctly configured.
	cmd := exec.Command("ssh", host, "-s", "sftp")
	cmd.Stderr = out

	wr, err := cmd.StdinPipe()
	if err != nil {
		return nil, nil, err
	}
	rd, err := cmd.StdoutPipe()
	if err != nil {
		return nil, nil, err
	}

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

	c, err := sftp.NewClientPipe(rd, wr)
	return c, func() error {
		return cmd.Wait()
	}, err
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"os/user"
	"path/filepath"
	"syscall"
	"time"

	"bazil.org/fuse"
	fusefs "bazil.org/fuse/fs"
	"bitbucket.org/anacrolix/go.torrent"
	"bitbucket.org/anacrolix/go.torrent/fs"
	metainfo "github.com/nsf/libtorgo/torrent"
)

var (
	downloadDir     string
	torrentPath     string
	mountDir        string
	disableTrackers = flag.Bool("disableTrackers", false, "disables trackers")
	testPeer        = flag.String("testPeer", "", "the address for a test peer")
	pprofAddr       = flag.String("pprofAddr", "", "pprof HTTP server bind address")
	testPeerAddr    *net.TCPAddr
)

func init() {
	flag.StringVar(&downloadDir, "downloadDir", "", "location to save torrent data")
	flag.StringVar(&torrentPath, "torrentPath", func() string {
		_user, err := user.Current()
		if err != nil {
			log.Fatal(err)
		}
		return filepath.Join(_user.HomeDir, ".config/transmission/torrents")
	}(), "torrent files in this location describe the contents of the mounted filesystem")
	flag.StringVar(&mountDir, "mountDir", "", "location the torrent contents are made available")
}

func resolveTestPeerAddr() {
	if *testPeer == "" {
		return
	}
	var err error
	testPeerAddr, err = net.ResolveTCPAddr("tcp4", *testPeer)
	if err != nil {
		log.Fatal(err)
	}
}

func setSignalHandlers() {
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		for {
			<-c
			err := fuse.Unmount(mountDir)
			if err != nil {
				log.Print(err)
			}
		}
	}()
}

func main() {
	flag.Parse()
	if flag.NArg() != 0 {
		os.Stderr.WriteString("one does not simply pass positional args\n")
		os.Exit(2)
	}
	if mountDir == "" {
		os.Stderr.WriteString("y u no specify mountpoint?\n")
		os.Exit(2)
	}
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	if *pprofAddr != "" {
		go http.ListenAndServe(*pprofAddr, nil)
	}
	conn, err := fuse.Mount(mountDir)
	if err != nil {
		log.Fatal(err)
	}
	defer fuse.Unmount(mountDir)
	// TODO: Think about the ramifications of exiting not due to a signal.
	setSignalHandlers()
	defer conn.Close()
	client := &torrent.Client{
		DataDir:         downloadDir,
		DisableTrackers: *disableTrackers,
	}
	client.Start()
	torrentDir, err := os.Open(torrentPath)
	defer torrentDir.Close()
	if err != nil {
		log.Fatal(err)
	}
	names, err := torrentDir.Readdirnames(-1)
	if err != nil {
		log.Fatal(err)
	}
	resolveTestPeerAddr()
	for _, name := range names {
		metaInfo, err := metainfo.LoadFromFile(filepath.Join(torrentPath, name))
		if err != nil {
			log.Print(err)
		}
		err = client.AddTorrent(metaInfo)
		if err != nil {
			log.Print(err)
		}
	}
	fs := torrentfs.New(client)
	go func() {
		for {
		torrentLoop:
			for _, t := range client.Torrents() {
				client.Lock()
				for _, c := range t.Conns {
					if c.Socket.RemoteAddr().String() == testPeerAddr.String() {
						client.Unlock()
						continue torrentLoop
					}
				}
				client.Unlock()
				if testPeerAddr != nil {
					if err := client.AddPeers(t.InfoHash, []torrent.Peer{{
						IP:   testPeerAddr.IP,
						Port: testPeerAddr.Port,
					}}); err != nil {
						log.Print(err)
					}
				}
			}
			time.Sleep(10 * time.Second)
		}
	}()
	if err := fusefs.Serve(conn, fs); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Break up torrentfs main(), it's too large<commit_after>package main

import (
	"flag"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"os/user"
	"path/filepath"
	"syscall"
	"time"

	"bazil.org/fuse"
	fusefs "bazil.org/fuse/fs"
	"bitbucket.org/anacrolix/go.torrent"
	"bitbucket.org/anacrolix/go.torrent/fs"
	metainfo "github.com/nsf/libtorgo/torrent"
)

var (
	downloadDir     string
	torrentPath     string
	mountDir        string
	disableTrackers = flag.Bool("disableTrackers", false, "disables trackers")
	testPeer        = flag.String("testPeer", "", "the address for a test peer")
	pprofAddr       = flag.String("pprofAddr", "", "pprof HTTP server bind address")
	testPeerAddr    *net.TCPAddr
)

func init() {
	flag.StringVar(&downloadDir, "downloadDir", "", "location to save torrent data")
	flag.StringVar(&torrentPath, "torrentPath", func() string {
		_user, err := user.Current()
		if err != nil {
			log.Fatal(err)
		}
		return filepath.Join(_user.HomeDir, ".config/transmission/torrents")
	}(), "torrent files in this location describe the contents of the mounted filesystem")
	flag.StringVar(&mountDir, "mountDir", "", "location the torrent contents are made available")
}

func resolveTestPeerAddr() {
	if *testPeer == "" {
		return
	}
	var err error
	testPeerAddr, err = net.ResolveTCPAddr("tcp4", *testPeer)
	if err != nil {
		log.Fatal(err)
	}
}

func setSignalHandlers() {
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		for {
			<-c
			err := fuse.Unmount(mountDir)
			if err != nil {
				log.Print(err)
			}
		}
	}()
}

func addTorrent(c *torrent.Client, file string) {
	metaInfo, err := metainfo.LoadFromFile(file)
	if err != nil {
		log.Print(err)
		return
	}
	err = c.AddTorrent(metaInfo)
	if err != nil {
		log.Print(err)
		return
	}
}

func addTorrentDir(c *torrent.Client, _path string) {
	torrentDir, err := os.Open(torrentPath)
	defer torrentDir.Close()
	if err != nil {
		log.Fatal(err)
	}
	names, err := torrentDir.Readdirnames(-1)
	if err != nil {
		log.Fatal(err)
	}
	for _, name := range names {
		go addTorrent(c, filepath.Join(_path, name))
	}
}

func addTestPeer(client *torrent.Client) {
torrents:
	for _, t := range client.Torrents() {
		client.Lock()
		for _, c := range t.Conns {
			if c.Socket.RemoteAddr().String() == testPeerAddr.String() {
				client.Unlock()
				continue torrents
			}
		}
		client.Unlock()
		if testPeerAddr != nil {
			if err := client.AddPeers(t.InfoHash, []torrent.Peer{{
				IP:   testPeerAddr.IP,
				Port: testPeerAddr.Port,
			}}); err != nil {
				log.Print(err)
			}
		}
	}
}

func main() {
	flag.Parse()
	if flag.NArg() != 0 {
		os.Stderr.WriteString("one does not simply pass positional args\n")
		os.Exit(2)
	}
	if mountDir == "" {
		os.Stderr.WriteString("y u no specify mountpoint?\n")
		os.Exit(2)
	}
	log.SetFlags(log.LstdFlags | log.Lshortfile)
	if *pprofAddr != "" {
		go http.ListenAndServe(*pprofAddr, nil)
	}
	conn, err := fuse.Mount(mountDir)
	if err != nil {
		log.Fatal(err)
	}
	defer fuse.Unmount(mountDir)
	// TODO: Think about the ramifications of exiting not due to a signal.
	setSignalHandlers()
	defer conn.Close()
	client := &torrent.Client{
		DataDir:         downloadDir,
		DisableTrackers: *disableTrackers,
	}
	client.Start()
	addTorrentDir(client, torrentPath)
	resolveTestPeerAddr()
	fs := torrentfs.New(client)
	go func() {
		for {
			addTestPeer(client)
			time.Sleep(10 * time.Second)
		}
	}()

	if err := fusefs.Serve(conn, fs); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


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

// Package etcd holds an etcd-specific implementation of the
// util.MasterElection interface.
package etcd

import (
	"context"
	"fmt"
	"strings"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/clientv3/concurrency"
	"github.com/golang/glog"
	"github.com/google/trillian/util"
)

// MasterElection is an implementation of util.MasterElection based on etcd.
type MasterElection struct {
	instanceID string
	treeID     int64
	lockFile   string
	client     *clientv3.Client
	session    *concurrency.Session
	election   *concurrency.Election
}

// Start commences election operation.
func (eme *MasterElection) Start(ctx context.Context) error {
	return nil
}

// WaitForMastership blocks until the current instance is master.
func (eme *MasterElection) WaitForMastership(ctx context.Context) error {
	return eme.election.Campaign(ctx, eme.instanceID)
}

// IsMaster returns whether the current instance is the master.
func (eme *MasterElection) IsMaster(ctx context.Context) (bool, error) {
	leader, err := eme.election.Leader(ctx)
	if err != nil {
		return false, err
	}
	return leader == eme.instanceID, nil
}

// ResignAndRestart releases mastership, and re-joins the election.
func (eme *MasterElection) ResignAndRestart(ctx context.Context) error {
	return eme.election.Resign(ctx)
}

// Close terminates election operation.
func (eme *MasterElection) Close(ctx context.Context) error {
	_ = eme.ResignAndRestart(ctx)
	if err := eme.session.Close(); err != nil {
		glog.Errorf("error closing session: %v", err)
	}
	return eme.client.Close()
}

// ElectionFactory creates etcd.MasterElection instances.
type ElectionFactory struct {
	client     *clientv3.Client
	instanceID string
	lockDir    string
}

// NewElectionFactory builds an election factory that uses the given parameters.
func NewElectionFactory(instanceID string, client *clientv3.Client, lockDir string) *ElectionFactory {
	return &ElectionFactory{
		client:     client,
		instanceID: instanceID,
		lockDir:    lockDir,
	}
}

// NewElection creates a specific etcd.MasterElection instance.
func (ef ElectionFactory) NewElection(ctx context.Context, treeID int64) (util.MasterElection, error) {
	session, err := concurrency.NewSession(ef.client)
	if err != nil {
		return nil, fmt.Errorf("failed to create etcd session: %v", err)
	}
	lockFile := fmt.Sprintf("%s/%d", strings.TrimRight(ef.lockDir, "/"), treeID)
	election := concurrency.NewElection(session, lockFile)

	eme := MasterElection{
		instanceID: ef.instanceID,
		treeID:     treeID,
		lockFile:   lockFile,
		client:     ef.client,
		session:    session,
		election:   election,
	}
	glog.Infof("MasterElection created: %+v", eme)
	return &eme, nil
}
<commit_msg>Update etcd election to new API<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package etcd holds an etcd-specific implementation of the
// util.MasterElection interface.
package etcd

import (
	"context"
	"fmt"
	"strings"

	"github.com/coreos/etcd/clientv3"
	"github.com/coreos/etcd/clientv3/concurrency"
	"github.com/golang/glog"
	"github.com/google/trillian/util"
)

// MasterElection is an implementation of util.MasterElection based on etcd.
type MasterElection struct {
	instanceID string
	treeID     int64
	lockFile   string
	client     *clientv3.Client
	session    *concurrency.Session
	election   *concurrency.Election
}

// Start commences election operation.
func (eme *MasterElection) Start(ctx context.Context) error {
	return nil
}

// WaitForMastership blocks until the current instance is master.
func (eme *MasterElection) WaitForMastership(ctx context.Context) error {
	return eme.election.Campaign(ctx, eme.instanceID)
}

// IsMaster returns whether the current instance is the master.
func (eme *MasterElection) IsMaster(ctx context.Context) (bool, error) {
	leader, err := eme.election.Leader(ctx)
	if err != nil {
		return false, err
	}
	return string(leader.Kvs[0].Value) == eme.instanceID, nil
}

// ResignAndRestart releases mastership, and re-joins the election.
func (eme *MasterElection) ResignAndRestart(ctx context.Context) error {
	return eme.election.Resign(ctx)
}

// Close terminates election operation.
func (eme *MasterElection) Close(ctx context.Context) error {
	_ = eme.ResignAndRestart(ctx)
	if err := eme.session.Close(); err != nil {
		glog.Errorf("error closing session: %v", err)
	}
	return eme.client.Close()
}

// ElectionFactory creates etcd.MasterElection instances.
type ElectionFactory struct {
	client     *clientv3.Client
	instanceID string
	lockDir    string
}

// NewElectionFactory builds an election factory that uses the given parameters.
func NewElectionFactory(instanceID string, client *clientv3.Client, lockDir string) *ElectionFactory {
	return &ElectionFactory{
		client:     client,
		instanceID: instanceID,
		lockDir:    lockDir,
	}
}

// NewElection creates a specific etcd.MasterElection instance.
func (ef ElectionFactory) NewElection(ctx context.Context, treeID int64) (util.MasterElection, error) {
	session, err := concurrency.NewSession(ef.client)
	if err != nil {
		return nil, fmt.Errorf("failed to create etcd session: %v", err)
	}
	lockFile := fmt.Sprintf("%s/%d", strings.TrimRight(ef.lockDir, "/"), treeID)
	election := concurrency.NewElection(session, lockFile)

	eme := MasterElection{
		instanceID: ef.instanceID,
		treeID:     treeID,
		lockFile:   lockFile,
		client:     ef.client,
		session:    session,
		election:   election,
	}
	glog.Infof("MasterElection created: %+v", eme)
	return &eme, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package events

import (
	"archive/tar"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/lib/session"
	"github.com/gravitational/teleport/lib/utils"

	"github.com/gravitational/trace"
	log "github.com/sirupsen/logrus"
)

// Header returns information about playback
type Header struct {
	// Tar detected tar format
	Tar bool
	// Proto is for proto format
	Proto bool
	// ProtoVersion is a version of the format, valid if Proto is true
	ProtoVersion int64
}

// DetectFormat detects format by reading first bytes
// of the header. Callers should call Seek()
// to reuse reader after calling this function.
func DetectFormat(r io.ReadSeeker) (*Header, error) {
	version := make([]byte, Int64Size)
	_, err := io.ReadFull(r, version)
	if err != nil {
		return nil, trace.ConvertSystemError(err)
	}
	protocolVersion := binary.BigEndian.Uint64(version)
	if protocolVersion == ProtoStreamV1 {
		return &Header{
			Proto:        true,
			ProtoVersion: int64(protocolVersion),
		}, nil
	}
	_, err = r.Seek(0, 0)
	if err != nil {
		return nil, trace.ConvertSystemError(err)
	}
	tr := tar.NewReader(r)
	_, err = tr.Next()
	if err != nil {
		return nil, trace.ConvertSystemError(err)
	}
	return &Header{Tar: true}, nil
}

// Export exports session file events to json, text or yaml
func Export(ctx context.Context, rs io.ReadSeeker, w io.Writer, exportFormat string) error {
	switch exportFormat {
	case teleport.Text:
	case teleport.JSON:
	default:
		return trace.BadParameter("unsupported format %q, %q is the only supported format", exportFormat, teleport.JSON)
	}

	format, err := DetectFormat(rs)
	if err != nil {
		return trace.Wrap(err)
	}
	_, err = rs.Seek(0, 0)
	if err != nil {
		return trace.ConvertSystemError(err)
	}
	switch {
	case format.Proto == true:
		protoReader := NewProtoReader(rs)
		for {
			event, err := protoReader.Read(ctx)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return trace.Wrap(err)
			}
			switch exportFormat {
			case teleport.JSON:
				data, err := utils.FastMarshal(event)
				if err != nil {
					return trace.ConvertSystemError(err)
				}
				_, err = fmt.Fprintln(w, string(data))
				if err != nil {
					return trace.ConvertSystemError(err)
				}
			default:
				return trace.BadParameter("unsupported format %q, %q is the only supported format", exportFormat, teleport.JSON)
			}
		}

	case format.Tar == true:
		return trace.BadParameter(
			"to review the events in format of teleport before version 4.4, extract the tarball and look inside")
	default:
		return trace.BadParameter("usupported format %v", format)
	}
}

// WriteForPlayback reads events from audit reader
// and writes them to the format optimized for playback
func WriteForPlayback(ctx context.Context, sid session.ID, reader AuditReader, dir string) error {
	w := &PlaybackWriter{
		sid:        sid,
		reader:     reader,
		dir:        dir,
		eventIndex: -1,
	}
	defer func() {
		if err := w.Close(); err != nil {
			log.WithError(err).Warningf("Failed to close writer.")
		}
	}()
	return w.Write(ctx)
}

// PlaybackWriter reads messages until end of file
// and writes them to directory in compatibility playback format
type PlaybackWriter struct {
	sid        session.ID
	dir        string
	reader     AuditReader
	indexFile  *os.File
	eventsFile *gzipWriter
	chunksFile *gzipWriter
	eventIndex int64
}

// Close closes all files
func (w *PlaybackWriter) Close() error {
	if w.indexFile != nil {
		w.indexFile.Close()
	}

	if w.chunksFile != nil {
		if err := w.chunksFile.Flush(); err != nil {
			log.Warningf("Failed to flush chunks file: %v.", err)
		}

		if err := w.chunksFile.Close(); err != nil {
			log.Warningf("Failed closing chunks file: %v.", err)
		}
	}

	if w.eventsFile != nil {
		if err := w.eventsFile.Flush(); err != nil {
			log.Warningf("Failed to flush events file: %v.", err)
		}

		if err := w.eventsFile.Close(); err != nil {
			log.Warningf("Failed closing events file: %v.", err)
		}
	}

	return nil
}

// Write writes the files in the format optimized for playback
func (w *PlaybackWriter) Write(ctx context.Context) error {
	if err := w.openIndexFile(); err != nil {
		return trace.Wrap(err)
	}
	for {
		event, err := w.reader.Read(ctx)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return trace.Wrap(err)
		}
		if err := w.writeEvent(event); err != nil {
			return trace.Wrap(err)
		}
	}
}

func (w *PlaybackWriter) writeEvent(event AuditEvent) error {
	switch event.GetType() {
	// Timing events for TTY playback go to both a chunks file (the raw bytes) as
	// well as well as the events file (structured events).
	case SessionPrintEvent:
		return trace.Wrap(w.writeSessionPrintEvent(event))
		// Playback does not use enhanced events at the moment,
		// so they are skipped
	case SessionCommandEvent, SessionDiskEvent, SessionNetworkEvent:
		return nil
	// All other events get put into the general events file. These are events like
	// session.join, session.end, etc.
	default:
		return trace.Wrap(w.writeRegularEvent(event))
	}
}

func (w *PlaybackWriter) writeSessionPrintEvent(event AuditEvent) error {
	print, ok := event.(*SessionPrint)
	if !ok {
		return trace.BadParameter("expected session print event, got %T", event)
	}
	w.eventIndex++
	event.SetIndex(w.eventIndex)
	if err := w.openEventsFile(0); err != nil {
		return trace.Wrap(err)
	}
	if err := w.openChunksFile(0); err != nil {
		return trace.Wrap(err)
	}
	data := print.Data
	print.Data = nil
	bytes, err := utils.FastMarshal(event)
	if err != nil {
		return trace.Wrap(err)
	}
	_, err = w.eventsFile.Write(append(bytes, '\n'))
	if err != nil {
		return trace.Wrap(err)
	}
	_, err = w.chunksFile.Write(data)
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

func (w *PlaybackWriter) writeRegularEvent(event AuditEvent) error {
	w.eventIndex++
	event.SetIndex(w.eventIndex)
	if err := w.openEventsFile(0); err != nil {
		return trace.Wrap(err)
	}
	bytes, err := utils.FastMarshal(event)
	if err != nil {
		return trace.Wrap(err)
	}
	_, err = w.eventsFile.Write(append(bytes, '\n'))
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

func (w *PlaybackWriter) openIndexFile() error {
	if w.indexFile != nil {
		return nil
	}
	var err error
	w.indexFile, err = os.OpenFile(
		filepath.Join(w.dir, fmt.Sprintf("%v.index", w.sid.String())), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0640)
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

func (w *PlaybackWriter) openEventsFile(eventIndex int64) error {
	if w.eventsFile != nil {
		return nil
	}
	eventsFileName := eventsFileName(w.dir, w.sid, "", eventIndex)

	// update the index file to write down that new events file has been created
	data, err := utils.FastMarshal(indexEntry{
		FileName: filepath.Base(eventsFileName),
		Type:     fileTypeEvents,
		Index:    eventIndex,
	})
	if err != nil {
		return trace.Wrap(err)
	}

	_, err = fmt.Fprintf(w.indexFile, "%v\n", string(data))
	if err != nil {
		return trace.Wrap(err)
	}

	// open new events file for writing
	file, err := os.OpenFile(eventsFileName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0640)
	if err != nil {
		return trace.Wrap(err)
	}
	w.eventsFile = newGzipWriter(file)
	return nil
}

func (w *PlaybackWriter) openChunksFile(offset int64) error {
	if w.chunksFile != nil {
		return nil
	}
	chunksFileName := chunksFileName(w.dir, w.sid, offset)

	// Update the index file to write down that new chunks file has been created.
	data, err := utils.FastMarshal(indexEntry{
		FileName: filepath.Base(chunksFileName),
		Type:     fileTypeChunks,
		Offset:   offset,
	})
	if err != nil {
		return trace.Wrap(err)
	}

	// index file will contain file name with extension .gz (assuming it was gzipped)
	_, err = fmt.Fprintf(w.indexFile, "%v\n", string(data))
	if err != nil {
		return trace.Wrap(err)
	}

	// open the chunks file for writing, but because the file is written without
	// compression, remove the .gz
	file, err := os.OpenFile(chunksFileName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0640)
	if err != nil {
		return trace.Wrap(err)
	}
	w.chunksFile = newGzipWriter(file)
	return nil
}
<commit_msg>Removed logic to export to text format.<commit_after>/*
Copyright 2020 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package events

import (
	"archive/tar"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/lib/session"
	"github.com/gravitational/teleport/lib/utils"

	"github.com/gravitational/trace"
	log "github.com/sirupsen/logrus"
)

// Header returns information about playback
type Header struct {
	// Tar detected tar format
	Tar bool
	// Proto is for proto format
	Proto bool
	// ProtoVersion is a version of the format, valid if Proto is true
	ProtoVersion int64
}

// DetectFormat detects format by reading first bytes
// of the header. Callers should call Seek()
// to reuse reader after calling this function.
func DetectFormat(r io.ReadSeeker) (*Header, error) {
	version := make([]byte, Int64Size)
	_, err := io.ReadFull(r, version)
	if err != nil {
		return nil, trace.ConvertSystemError(err)
	}
	protocolVersion := binary.BigEndian.Uint64(version)
	if protocolVersion == ProtoStreamV1 {
		return &Header{
			Proto:        true,
			ProtoVersion: int64(protocolVersion),
		}, nil
	}
	_, err = r.Seek(0, 0)
	if err != nil {
		return nil, trace.ConvertSystemError(err)
	}
	tr := tar.NewReader(r)
	_, err = tr.Next()
	if err != nil {
		return nil, trace.ConvertSystemError(err)
	}
	return &Header{Tar: true}, nil
}

// Export converts session files from binary/protobuf to text/JSON.
func Export(ctx context.Context, rs io.ReadSeeker, w io.Writer, exportFormat string) error {
	switch exportFormat {
	case teleport.JSON:
	default:
		return trace.BadParameter("unsupported format %q, %q is the only supported format", exportFormat, teleport.JSON)
	}

	format, err := DetectFormat(rs)
	if err != nil {
		return trace.Wrap(err)
	}
	_, err = rs.Seek(0, 0)
	if err != nil {
		return trace.ConvertSystemError(err)
	}
	switch {
	case format.Proto == true:
		protoReader := NewProtoReader(rs)
		for {
			event, err := protoReader.Read(ctx)
			if err != nil {
				if err == io.EOF {
					return nil
				}
				return trace.Wrap(err)
			}
			switch exportFormat {
			case teleport.JSON:
				data, err := utils.FastMarshal(event)
				if err != nil {
					return trace.ConvertSystemError(err)
				}
				_, err = fmt.Fprintln(w, string(data))
				if err != nil {
					return trace.ConvertSystemError(err)
				}
			default:
				return trace.BadParameter("unsupported format %q, %q is the only supported format", exportFormat, teleport.JSON)
			}
		}
	case format.Tar == true:
		return trace.BadParameter(
			"to review the events in format of teleport before version 4.4, extract the tarball and look inside")
	default:
		return trace.BadParameter("unsupported format %v", format)
	}
}

// WriteForPlayback reads events from audit reader
// and writes them to the format optimized for playback
func WriteForPlayback(ctx context.Context, sid session.ID, reader AuditReader, dir string) error {
	w := &PlaybackWriter{
		sid:        sid,
		reader:     reader,
		dir:        dir,
		eventIndex: -1,
	}
	defer func() {
		if err := w.Close(); err != nil {
			log.WithError(err).Warningf("Failed to close writer.")
		}
	}()
	return w.Write(ctx)
}

// PlaybackWriter reads messages until end of file
// and writes them to directory in compatibility playback format
type PlaybackWriter struct {
	sid        session.ID
	dir        string
	reader     AuditReader
	indexFile  *os.File
	eventsFile *gzipWriter
	chunksFile *gzipWriter
	eventIndex int64
}

// Close closes all files
func (w *PlaybackWriter) Close() error {
	if w.indexFile != nil {
		w.indexFile.Close()
	}

	if w.chunksFile != nil {
		if err := w.chunksFile.Flush(); err != nil {
			log.Warningf("Failed to flush chunks file: %v.", err)
		}

		if err := w.chunksFile.Close(); err != nil {
			log.Warningf("Failed closing chunks file: %v.", err)
		}
	}

	if w.eventsFile != nil {
		if err := w.eventsFile.Flush(); err != nil {
			log.Warningf("Failed to flush events file: %v.", err)
		}

		if err := w.eventsFile.Close(); err != nil {
			log.Warningf("Failed closing events file: %v.", err)
		}
	}

	return nil
}

// Write writes the files in the format optimized for playback
func (w *PlaybackWriter) Write(ctx context.Context) error {
	if err := w.openIndexFile(); err != nil {
		return trace.Wrap(err)
	}
	for {
		event, err := w.reader.Read(ctx)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return trace.Wrap(err)
		}
		if err := w.writeEvent(event); err != nil {
			return trace.Wrap(err)
		}
	}
}

func (w *PlaybackWriter) writeEvent(event AuditEvent) error {
	switch event.GetType() {
	// Timing events for TTY playback go to both a chunks file (the raw bytes) as
	// well as well as the events file (structured events).
	case SessionPrintEvent:
		return trace.Wrap(w.writeSessionPrintEvent(event))
		// Playback does not use enhanced events at the moment,
		// so they are skipped
	case SessionCommandEvent, SessionDiskEvent, SessionNetworkEvent:
		return nil
	// All other events get put into the general events file. These are events like
	// session.join, session.end, etc.
	default:
		return trace.Wrap(w.writeRegularEvent(event))
	}
}

func (w *PlaybackWriter) writeSessionPrintEvent(event AuditEvent) error {
	print, ok := event.(*SessionPrint)
	if !ok {
		return trace.BadParameter("expected session print event, got %T", event)
	}
	w.eventIndex++
	event.SetIndex(w.eventIndex)
	if err := w.openEventsFile(0); err != nil {
		return trace.Wrap(err)
	}
	if err := w.openChunksFile(0); err != nil {
		return trace.Wrap(err)
	}
	data := print.Data
	print.Data = nil
	bytes, err := utils.FastMarshal(event)
	if err != nil {
		return trace.Wrap(err)
	}
	_, err = w.eventsFile.Write(append(bytes, '\n'))
	if err != nil {
		return trace.Wrap(err)
	}
	_, err = w.chunksFile.Write(data)
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

func (w *PlaybackWriter) writeRegularEvent(event AuditEvent) error {
	w.eventIndex++
	event.SetIndex(w.eventIndex)
	if err := w.openEventsFile(0); err != nil {
		return trace.Wrap(err)
	}
	bytes, err := utils.FastMarshal(event)
	if err != nil {
		return trace.Wrap(err)
	}
	_, err = w.eventsFile.Write(append(bytes, '\n'))
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

func (w *PlaybackWriter) openIndexFile() error {
	if w.indexFile != nil {
		return nil
	}
	var err error
	w.indexFile, err = os.OpenFile(
		filepath.Join(w.dir, fmt.Sprintf("%v.index", w.sid.String())), os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0640)
	if err != nil {
		return trace.Wrap(err)
	}
	return nil
}

func (w *PlaybackWriter) openEventsFile(eventIndex int64) error {
	if w.eventsFile != nil {
		return nil
	}
	eventsFileName := eventsFileName(w.dir, w.sid, "", eventIndex)

	// update the index file to write down that new events file has been created
	data, err := utils.FastMarshal(indexEntry{
		FileName: filepath.Base(eventsFileName),
		Type:     fileTypeEvents,
		Index:    eventIndex,
	})
	if err != nil {
		return trace.Wrap(err)
	}

	_, err = fmt.Fprintf(w.indexFile, "%v\n", string(data))
	if err != nil {
		return trace.Wrap(err)
	}

	// open new events file for writing
	file, err := os.OpenFile(eventsFileName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0640)
	if err != nil {
		return trace.Wrap(err)
	}
	w.eventsFile = newGzipWriter(file)
	return nil
}

func (w *PlaybackWriter) openChunksFile(offset int64) error {
	if w.chunksFile != nil {
		return nil
	}
	chunksFileName := chunksFileName(w.dir, w.sid, offset)

	// Update the index file to write down that new chunks file has been created.
	data, err := utils.FastMarshal(indexEntry{
		FileName: filepath.Base(chunksFileName),
		Type:     fileTypeChunks,
		Offset:   offset,
	})
	if err != nil {
		return trace.Wrap(err)
	}

	// index file will contain file name with extension .gz (assuming it was gzipped)
	_, err = fmt.Fprintf(w.indexFile, "%v\n", string(data))
	if err != nil {
		return trace.Wrap(err)
	}

	// open the chunks file for writing, but because the file is written without
	// compression, remove the .gz
	file, err := os.OpenFile(chunksFileName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0640)
	if err != nil {
		return trace.Wrap(err)
	}
	w.chunksFile = newGzipWriter(file)
	return nil
}
<|endoftext|>


<commit_before>package e2etest

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/hashicorp/terraform/internal/e2e"
	"github.com/hashicorp/terraform/internal/getproviders"
)

// TestProviderTampering tests various ways that the provider plugins in the
// local cache directory might be modified after an initial "terraform init",
// which other Terraform commands which use those plugins should catch and
// report early.
func TestProviderTampering(t *testing.T) {
	// General setup: we'll do a one-off init of a test directory as our
	// starting point, and then we'll clone that result for each test so
	// that we can save the cost of a repeated re-init with the same
	// provider.
	t.Parallel()

	// This test reaches out to releases.hashicorp.com to download the
	// null provider, so it can only run if network access is allowed.
	skipIfCannotAccessNetwork(t)

	fixturePath := filepath.Join("testdata", "provider-tampering-base")
	tf := e2e.NewBinary(terraformBin, fixturePath)
	defer tf.Close()

	stdout, stderr, err := tf.Run("init")
	if err != nil {
		t.Fatalf("unexpected init error: %s\nstderr:\n%s", err, stderr)
	}
	if !strings.Contains(stdout, "Installing hashicorp/null v") {
		t.Errorf("null provider download message is missing from init output:\n%s", stdout)
		t.Logf("(this can happen if you have a copy of the plugin in one of the global plugin search dirs)")
	}

	seedDir := tf.WorkDir()
	const providerVersion = "3.1.0" // must match the version in the fixture config
	pluginDir := ".terraform/providers/registry.terraform.io/hashicorp/null/" + providerVersion + "/" + getproviders.CurrentPlatform.String()
	pluginExe := pluginDir + "/terraform-provider-null_v" + providerVersion + "_x5"
	if getproviders.CurrentPlatform.OS == "windows" {
		pluginExe += ".exe" // ugh
	}

	t.Run("cache dir totally gone", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		err := os.RemoveAll(filepath.Join(workDir, ".terraform"))
		if err != nil {
			t.Fatal(err)
		}

		_, stderr, err := tf.Run("plan")
		if err == nil {
			t.Fatalf("unexpected plan success\nstdout:\n%s", stdout)
		}
		if want := `registry.terraform.io/hashicorp/null: there is no package for registry.terraform.io/hashicorp/null 3.1.0 cached in .terraform/providers`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `terraform init`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("null plugin package modified before plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		err := ioutil.WriteFile(filepath.Join(workDir, pluginExe), []byte("tamper"), 0600)
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("plan")
		if err == nil {
			t.Fatalf("unexpected plan success\nstdout:\n%s", stdout)
		}
		if want := `registry.terraform.io/hashicorp/null: the cached package for registry.terraform.io/hashicorp/null 3.1.0 (in .terraform/providers) does not match any of the checksums recorded in the dependency lock file`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `terraform init`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("version constraint changed in config before plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		err := ioutil.WriteFile(filepath.Join(workDir, "provider-tampering-base.tf"), []byte(`
			terraform {
				required_providers {
					null = {
						source  = "hashicorp/null"
						version = "1.0.0"
					}
				}
			}
		`), 0600)
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("plan")
		if err == nil {
			t.Fatalf("unexpected plan success\nstdout:\n%s", stdout)
		}
		if want := `provider registry.terraform.io/hashicorp/null: locked version selection 3.1.0 doesn't match the updated version constraints "1.0.0"`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `terraform init -upgrade`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("lock file modified before plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		// NOTE: We're just emptying out the lock file here because that's
		// good enough for what we're trying to assert. The leaf codepath
		// that generates this family of errors has some different variations
		// of this error message for otehr sorts of inconsistency, but those
		// are tested more thoroughly over in the "configs" package, which is
		// ultimately responsible for that logic.
		err := ioutil.WriteFile(filepath.Join(workDir, ".terraform.lock.hcl"), []byte(``), 0600)
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("plan")
		if err == nil {
			t.Fatalf("unexpected plan success\nstdout:\n%s", stdout)
		}
		if want := `provider registry.terraform.io/hashicorp/null: required by this configuration but no version is selected`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `terraform init`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("lock file modified after plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		_, stderr, err := tf.Run("plan", "-out", "tfplan")
		if err != nil {
			t.Fatalf("unexpected plan failure\nstderr:\n%s", stderr)
		}

		err = os.Remove(filepath.Join(workDir, ".terraform.lock.hcl"))
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("apply", "tfplan")
		if err == nil {
			t.Fatalf("unexpected apply success\nstdout:\n%s", stdout)
		}
		if want := `provider registry.terraform.io/hashicorp/null: required by this configuration but no version is selected`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `Create a new plan from the updated configuration.`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("plugin cache dir entirely removed after plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		_, stderr, err := tf.Run("plan", "-out", "tfplan")
		if err != nil {
			t.Fatalf("unexpected plan failure\nstderr:\n%s", stderr)
		}

		err = os.RemoveAll(filepath.Join(workDir, ".terraform"))
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("apply", "tfplan")
		if err == nil {
			t.Fatalf("unexpected apply success\nstdout:\n%s", stdout)
		}
		if want := `registry.terraform.io/hashicorp/null: there is no package for registry.terraform.io/hashicorp/null 3.1.0 cached in .terraform/providers`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("null plugin package modified after plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		_, stderr, err := tf.Run("plan", "-out", "tfplan")
		if err != nil {
			t.Fatalf("unexpected plan failure\nstderr:\n%s", stderr)
		}

		err = ioutil.WriteFile(filepath.Join(workDir, pluginExe), []byte("tamper"), 0600)
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("apply", "tfplan")
		if err == nil {
			t.Fatalf("unexpected apply success\nstdout:\n%s", stdout)
		}
		if want := `registry.terraform.io/hashicorp/null: the cached package for registry.terraform.io/hashicorp/null 3.1.0 (in .terraform/providers) does not match any of the checksums recorded in the dependency lock file`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
}
<commit_msg>command/e2etest: Ensure init fixes missing cache<commit_after>package e2etest

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/hashicorp/terraform/internal/e2e"
	"github.com/hashicorp/terraform/internal/getproviders"
)

// TestProviderTampering tests various ways that the provider plugins in the
// local cache directory might be modified after an initial "terraform init",
// which other Terraform commands which use those plugins should catch and
// report early.
func TestProviderTampering(t *testing.T) {
	// General setup: we'll do a one-off init of a test directory as our
	// starting point, and then we'll clone that result for each test so
	// that we can save the cost of a repeated re-init with the same
	// provider.
	t.Parallel()

	// This test reaches out to releases.hashicorp.com to download the
	// null provider, so it can only run if network access is allowed.
	skipIfCannotAccessNetwork(t)

	fixturePath := filepath.Join("testdata", "provider-tampering-base")
	tf := e2e.NewBinary(terraformBin, fixturePath)
	defer tf.Close()

	stdout, stderr, err := tf.Run("init")
	if err != nil {
		t.Fatalf("unexpected init error: %s\nstderr:\n%s", err, stderr)
	}
	if !strings.Contains(stdout, "Installing hashicorp/null v") {
		t.Errorf("null provider download message is missing from init output:\n%s", stdout)
		t.Logf("(this can happen if you have a copy of the plugin in one of the global plugin search dirs)")
	}

	seedDir := tf.WorkDir()
	const providerVersion = "3.1.0" // must match the version in the fixture config
	pluginDir := ".terraform/providers/registry.terraform.io/hashicorp/null/" + providerVersion + "/" + getproviders.CurrentPlatform.String()
	pluginExe := pluginDir + "/terraform-provider-null_v" + providerVersion + "_x5"
	if getproviders.CurrentPlatform.OS == "windows" {
		pluginExe += ".exe" // ugh
	}

	t.Run("cache dir totally gone", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		err := os.RemoveAll(filepath.Join(workDir, ".terraform"))
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("plan")
		if err == nil {
			t.Fatalf("unexpected plan success\nstdout:\n%s", stdout)
		}
		if want := `registry.terraform.io/hashicorp/null: there is no package for registry.terraform.io/hashicorp/null 3.1.0 cached in .terraform/providers`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `terraform init`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}

		// Running init as suggested resolves the problem
		_, stderr, err = tf.Run("init")
		if err != nil {
			t.Fatalf("unexpected init error: %s\nstderr:\n%s", err, stderr)
		}
		_, stderr, err = tf.Run("plan")
		if err != nil {
			t.Fatalf("unexpected plan error: %s\nstderr:\n%s", err, stderr)
		}
	})
	t.Run("null plugin package modified before plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		err := ioutil.WriteFile(filepath.Join(workDir, pluginExe), []byte("tamper"), 0600)
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("plan")
		if err == nil {
			t.Fatalf("unexpected plan success\nstdout:\n%s", stdout)
		}
		if want := `registry.terraform.io/hashicorp/null: the cached package for registry.terraform.io/hashicorp/null 3.1.0 (in .terraform/providers) does not match any of the checksums recorded in the dependency lock file`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `terraform init`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("version constraint changed in config before plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		err := ioutil.WriteFile(filepath.Join(workDir, "provider-tampering-base.tf"), []byte(`
			terraform {
				required_providers {
					null = {
						source  = "hashicorp/null"
						version = "1.0.0"
					}
				}
			}
		`), 0600)
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("plan")
		if err == nil {
			t.Fatalf("unexpected plan success\nstdout:\n%s", stdout)
		}
		if want := `provider registry.terraform.io/hashicorp/null: locked version selection 3.1.0 doesn't match the updated version constraints "1.0.0"`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `terraform init -upgrade`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("lock file modified before plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		// NOTE: We're just emptying out the lock file here because that's
		// good enough for what we're trying to assert. The leaf codepath
		// that generates this family of errors has some different variations
		// of this error message for otehr sorts of inconsistency, but those
		// are tested more thoroughly over in the "configs" package, which is
		// ultimately responsible for that logic.
		err := ioutil.WriteFile(filepath.Join(workDir, ".terraform.lock.hcl"), []byte(``), 0600)
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("plan")
		if err == nil {
			t.Fatalf("unexpected plan success\nstdout:\n%s", stdout)
		}
		if want := `provider registry.terraform.io/hashicorp/null: required by this configuration but no version is selected`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `terraform init`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("lock file modified after plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		_, stderr, err := tf.Run("plan", "-out", "tfplan")
		if err != nil {
			t.Fatalf("unexpected plan failure\nstderr:\n%s", stderr)
		}

		err = os.Remove(filepath.Join(workDir, ".terraform.lock.hcl"))
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("apply", "tfplan")
		if err == nil {
			t.Fatalf("unexpected apply success\nstdout:\n%s", stdout)
		}
		if want := `provider registry.terraform.io/hashicorp/null: required by this configuration but no version is selected`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
		if want := `Create a new plan from the updated configuration.`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("plugin cache dir entirely removed after plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		_, stderr, err := tf.Run("plan", "-out", "tfplan")
		if err != nil {
			t.Fatalf("unexpected plan failure\nstderr:\n%s", stderr)
		}

		err = os.RemoveAll(filepath.Join(workDir, ".terraform"))
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("apply", "tfplan")
		if err == nil {
			t.Fatalf("unexpected apply success\nstdout:\n%s", stdout)
		}
		if want := `registry.terraform.io/hashicorp/null: there is no package for registry.terraform.io/hashicorp/null 3.1.0 cached in .terraform/providers`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
	t.Run("null plugin package modified after plan", func(t *testing.T) {
		tf := e2e.NewBinary(terraformBin, seedDir)
		defer tf.Close()
		workDir := tf.WorkDir()

		_, stderr, err := tf.Run("plan", "-out", "tfplan")
		if err != nil {
			t.Fatalf("unexpected plan failure\nstderr:\n%s", stderr)
		}

		err = ioutil.WriteFile(filepath.Join(workDir, pluginExe), []byte("tamper"), 0600)
		if err != nil {
			t.Fatal(err)
		}

		stdout, stderr, err := tf.Run("apply", "tfplan")
		if err == nil {
			t.Fatalf("unexpected apply success\nstdout:\n%s", stdout)
		}
		if want := `registry.terraform.io/hashicorp/null: the cached package for registry.terraform.io/hashicorp/null 3.1.0 (in .terraform/providers) does not match any of the checksums recorded in the dependency lock file`; !strings.Contains(stderr, want) {
			t.Errorf("missing expected error message\nwant substring: %s\ngot:\n%s", want, stderr)
		}
	})
}
<|endoftext|>


<commit_before>package wats

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	. "github.com/cloudfoundry-incubator/cf-test-helpers/workflowhelpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
)

const (
	DEFAULT_TIMEOUT = 45 * time.Second
	CF_PUSH_TIMEOUT = 3 * time.Minute
)

var (
	appName            string
	config             *watsConfig
	environment        *ReproducibleTestSuiteSetup
	hwcBuildPackURL    = "https://github.com/cloudfoundry-incubator/hwc-buildpack/releases/download/v2.1.2/hwc_buildpack-cached-v2.1.2.zip"
	binaryBuildPackURL = "https://github.com/cloudfoundry/binary-buildpack/releases/download/v1.0.8/binary_buildpack-cached-v1.0.8.zip"
)

func guidForAppName(appName string) string {
	cfApp := cf.Cf("app", appName, "--guid")
	Expect(cfApp.Wait()).To(Exit(0))

	appGuid := strings.TrimSpace(string(cfApp.Out.Contents()))
	Expect(appGuid).NotTo(Equal(""))
	return appGuid
}

func guidForSpaceName(spaceName string) string {
	cfSpace := cf.Cf("space", spaceName, "--guid")
	Expect(cfSpace.Wait()).To(Exit(0))

	spaceGuid := strings.TrimSpace(string(cfSpace.Out.Contents()))
	Expect(spaceGuid).NotTo(Equal(""))
	return spaceGuid
}

func enableDiego(appName string) {
	guid := guidForAppName(appName)
	Eventually(cf.Cf("curl", "/v2/apps/"+guid, "-X", "PUT", "-d", `{"diego": true}`)).Should(Exit(0))
}

func disableHealthCheck(appName string) {
	guid := guidForAppName(appName)
	Eventually(cf.Cf("curl", "/v2/apps/"+guid, "-X", "PUT", "-d", `{"health_check_type":"none"}`)).Should(Exit(0))
}

func TestDiegoWindows(t *testing.T) {
	RegisterFailHandler(Fail)

	SetDefaultEventuallyTimeout(time.Minute)
	SetDefaultEventuallyPollingInterval(time.Second)

	var err error
	config, err = LoadWatsConfig()
	if err != nil {
		t.Fatalf("could not load WATS config", err)
	}

	if config.NumWindowsCells == 0 {
		t.Fatalf("Please provide 'num_windows_cells' as a property in the integration config JSON (The number of windows cells in tested deployment)")
	}

	environment = NewTestSuiteSetup(config)

	BeforeSuite(func() {
		environment.Setup()
		binaryBuildpackVersion := getBuildpackVersion("binary_buildpack")
		if versionGreaterThan(binaryBuildpackVersion, 1, 0, 7) {
			binaryBuildPackURL = "binary_buildpack"
		}

		hwcBuildpackVersion := getBuildpackVersion("hwc_buildpack")
		if versionGreaterThan(hwcBuildpackVersion, 2, 0, 0) {
			hwcBuildPackURL = "hwc_buildpack"
		}

		if config.GetIsolationSegmentName() != "" {
			AsUser(environment.AdminUserContext(), environment.ShortTimeout(), func() {
				isoSegGuid := createOrGetIsolationSegment(config.GetIsolationSegmentName())
				attachIsolationSegmentToOrg(environment, isoSegGuid)
				attachIsolationSegmentToSpace(environment, isoSegGuid)
			})
		}
	})

	AfterSuite(func() {
		environment.Teardown()
	})

	BeforeEach(func() {
		Eventually(cf.Cf("apps").Out).Should(Say("No apps found"))
		appName = generator.PrefixedRandomName(config.GetNamePrefix(), "APP")
	})

	AfterEach(func() {
		Eventually(cf.Cf("logs", appName, "--recent")).Should(Exit())
		Eventually(cf.Cf("delete", appName, "-f")).Should(Exit(0))
	})

	componentName := "DiegoWindows"

	rs := []Reporter{}

	if config.GetArtifactsDirectory() != "" {
		helpers.EnableCFTrace(config, componentName)
		rs = append(rs, helpers.NewJUnitReporter(config, componentName))
	}

	RunSpecsWithDefaultAndCustomReporters(t, componentName, rs)
}
func getBuildpackVersion(name string) string {
	buildpack := cf.Cf("curl", fmt.Sprintf("/v2/buildpacks?q=name:%s", name))
	Expect(buildpack.Wait()).To(Exit(0))
	type Buildpack struct {
		Resources []struct {
			Entity struct {
				FileName string
			}
		}
	}
	var b Buildpack
	Expect(json.Unmarshal(buildpack.Out.Contents(), &b)).To(Succeed())
	if len(b.Resources) == 0 {
		return "0.0.0"
	}
	re := regexp.MustCompile(`[0-9]+\.[0-9]+\.[0-9]+`)
	return re.FindString(b.Resources[0].Entity.FileName)
}

func versionGreaterThan(version string, inputMajor, inputMinor, inputPatch int) bool {
	versions := strings.Split(version, ".")
	major, _ := strconv.Atoi(versions[0])
	if major > inputMajor {
		return true
	}
	minor, _ := strconv.Atoi(versions[1])
	if minor > inputMinor {
		return true
	}
	patch, _ := strconv.Atoi(versions[2])
	if patch > inputPatch {
		return true
	}
	return false
}

func createOrGetIsolationSegment(isolationSegmentName string) string {
	// This could go in cf-test-helpers
	guid := getIsolationSegmentGuid(isolationSegmentName)
	if guid == "" {
		Eventually(cf.Cf("curl", "/v3/isolation_segments/", "-X", "POST", "-d", fmt.Sprintf(`{"name":"%s"}`, isolationSegmentName))).Should(Exit(0))
		guid = getIsolationSegmentGuid(isolationSegmentName)
	}
	return guid
}

func getIsolationSegmentGuid(isolationSegmentName string) string {
	return getV3ResourceGuid(fmt.Sprintf("/v3/isolation_segments?names=%s", isolationSegmentName))
}

func attachIsolationSegmentToOrg(environment *ReproducibleTestSuiteSetup, isoSegGuid string) {
	orgGuid := getOrganizationGuid(environment.GetOrganizationName())
	response := cf.Cf(
		"curl",
		fmt.Sprintf("/v3/isolation_segments/%s/relationships/organizations", isoSegGuid),
		"-X",
		"POST",
		"-d",
		fmt.Sprintf(`{"data":[{"guid": "%s"}]}`, orgGuid),
	)
	Expect(response.Wait()).To(Exit(0))
}

func attachIsolationSegmentToSpace(environment *ReproducibleTestSuiteSetup, isoSegGuid string) {
	spaceGuid := getSpaceGuidForOrg(getOrganizationGuid(environment.GetOrganizationName()))
	response := cf.Cf(
		"curl",
		fmt.Sprintf("/v2/spaces/%s", spaceGuid),
		"-X",
		"PUT",
		"-d",
		fmt.Sprintf(`{"isolation_segment_guid":"%s"}`, isoSegGuid),
	)
	Expect(response.Wait()).To(Exit(0))
}

func getOrganizationGuid(organizationName string) string {
	return getV2ResourceGuid(fmt.Sprintf("/v2/organizations?q=name:%s", organizationName))
}

func getSpaceGuidForOrg(orgGuid string) string {
	return getV2ResourceGuid(fmt.Sprintf("/v2/organizations/%s/spaces", orgGuid))
}

func getV2ResourceGuid(endpoint string) string {
	response := cf.Cf("curl", endpoint, "-X", "GET")
	Expect(response.Wait()).To(Exit(0))
	var r struct {
		Resources []struct {
			Metadata struct {
				Guid string
			}
		}
	}
	Expect(json.Unmarshal(response.Out.Contents(), &r)).To(Succeed())
	Expect(len(r.Resources)).To(Equal(1))
	return r.Resources[0].Metadata.Guid
}

func getV3ResourceGuid(endpoint string) string {
	response := cf.Cf("curl", endpoint, "-X", "GET")
	Expect(response.Wait()).To(Exit(0))
	var r struct {
		Resources []struct {
			Guid string
		}
	}
	Expect(json.Unmarshal(response.Out.Contents(), &r)).To(Succeed())
	if len(r.Resources) == 0 {
		return ""
	}
	return r.Resources[0].Guid
}
<commit_msg>Bump binary_buildpack to 1.0.10<commit_after>package wats

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

	. "github.com/cloudfoundry-incubator/cf-test-helpers/workflowhelpers"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	. "github.com/onsi/gomega/gbytes"
	. "github.com/onsi/gomega/gexec"

	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	"github.com/cloudfoundry-incubator/cf-test-helpers/generator"
	"github.com/cloudfoundry-incubator/cf-test-helpers/helpers"
)

const (
	DEFAULT_TIMEOUT = 45 * time.Second
	CF_PUSH_TIMEOUT = 3 * time.Minute
)

var (
	appName            string
	config             *watsConfig
	environment        *ReproducibleTestSuiteSetup
	hwcBuildPackURL    = "https://github.com/cloudfoundry-incubator/hwc-buildpack/releases/download/v2.1.2/hwc_buildpack-cached-v2.1.2.zip"
	binaryBuildPackURL = "https://github.com/cloudfoundry/binary-buildpack/releases/download/v1.0.10/binary_buildpack-cached-v1.0.10.zip"
)

func guidForAppName(appName string) string {
	cfApp := cf.Cf("app", appName, "--guid")
	Expect(cfApp.Wait()).To(Exit(0))

	appGuid := strings.TrimSpace(string(cfApp.Out.Contents()))
	Expect(appGuid).NotTo(Equal(""))
	return appGuid
}

func guidForSpaceName(spaceName string) string {
	cfSpace := cf.Cf("space", spaceName, "--guid")
	Expect(cfSpace.Wait()).To(Exit(0))

	spaceGuid := strings.TrimSpace(string(cfSpace.Out.Contents()))
	Expect(spaceGuid).NotTo(Equal(""))
	return spaceGuid
}

func enableDiego(appName string) {
	guid := guidForAppName(appName)
	Eventually(cf.Cf("curl", "/v2/apps/"+guid, "-X", "PUT", "-d", `{"diego": true}`)).Should(Exit(0))
}

func disableHealthCheck(appName string) {
	guid := guidForAppName(appName)
	Eventually(cf.Cf("curl", "/v2/apps/"+guid, "-X", "PUT", "-d", `{"health_check_type":"none"}`)).Should(Exit(0))
}

func TestDiegoWindows(t *testing.T) {
	RegisterFailHandler(Fail)

	SetDefaultEventuallyTimeout(time.Minute)
	SetDefaultEventuallyPollingInterval(time.Second)

	var err error
	config, err = LoadWatsConfig()
	if err != nil {
		t.Fatalf("could not load WATS config", err)
	}

	if config.NumWindowsCells == 0 {
		t.Fatalf("Please provide 'num_windows_cells' as a property in the integration config JSON (The number of windows cells in tested deployment)")
	}

	environment = NewTestSuiteSetup(config)

	BeforeSuite(func() {
		environment.Setup()
		binaryBuildpackVersion := getBuildpackVersion("binary_buildpack")
		if versionGreaterThan(binaryBuildpackVersion, 1, 0, 7) {
			binaryBuildPackURL = "binary_buildpack"
		}

		hwcBuildpackVersion := getBuildpackVersion("hwc_buildpack")
		if versionGreaterThan(hwcBuildpackVersion, 2, 0, 0) {
			hwcBuildPackURL = "hwc_buildpack"
		}

		if config.GetIsolationSegmentName() != "" {
			AsUser(environment.AdminUserContext(), environment.ShortTimeout(), func() {
				isoSegGuid := createOrGetIsolationSegment(config.GetIsolationSegmentName())
				attachIsolationSegmentToOrg(environment, isoSegGuid)
				attachIsolationSegmentToSpace(environment, isoSegGuid)
			})
		}
	})

	AfterSuite(func() {
		environment.Teardown()
	})

	BeforeEach(func() {
		Eventually(cf.Cf("apps").Out).Should(Say("No apps found"))
		appName = generator.PrefixedRandomName(config.GetNamePrefix(), "APP")
	})

	AfterEach(func() {
		Eventually(cf.Cf("logs", appName, "--recent")).Should(Exit())
		Eventually(cf.Cf("delete", appName, "-f")).Should(Exit(0))
	})

	componentName := "DiegoWindows"

	rs := []Reporter{}

	if config.GetArtifactsDirectory() != "" {
		helpers.EnableCFTrace(config, componentName)
		rs = append(rs, helpers.NewJUnitReporter(config, componentName))
	}

	RunSpecsWithDefaultAndCustomReporters(t, componentName, rs)
}
func getBuildpackVersion(name string) string {
	buildpack := cf.Cf("curl", fmt.Sprintf("/v2/buildpacks?q=name:%s", name))
	Expect(buildpack.Wait()).To(Exit(0))
	type Buildpack struct {
		Resources []struct {
			Entity struct {
				FileName string
			}
		}
	}
	var b Buildpack
	Expect(json.Unmarshal(buildpack.Out.Contents(), &b)).To(Succeed())
	if len(b.Resources) == 0 {
		return "0.0.0"
	}
	re := regexp.MustCompile(`[0-9]+\.[0-9]+\.[0-9]+`)
	return re.FindString(b.Resources[0].Entity.FileName)
}

func versionGreaterThan(version string, inputMajor, inputMinor, inputPatch int) bool {
	versions := strings.Split(version, ".")
	major, _ := strconv.Atoi(versions[0])
	if major > inputMajor {
		return true
	}
	minor, _ := strconv.Atoi(versions[1])
	if minor > inputMinor {
		return true
	}
	patch, _ := strconv.Atoi(versions[2])
	if patch > inputPatch {
		return true
	}
	return false
}

func createOrGetIsolationSegment(isolationSegmentName string) string {
	// This could go in cf-test-helpers
	guid := getIsolationSegmentGuid(isolationSegmentName)
	if guid == "" {
		Eventually(cf.Cf("curl", "/v3/isolation_segments/", "-X", "POST", "-d", fmt.Sprintf(`{"name":"%s"}`, isolationSegmentName))).Should(Exit(0))
		guid = getIsolationSegmentGuid(isolationSegmentName)
	}
	return guid
}

func getIsolationSegmentGuid(isolationSegmentName string) string {
	return getV3ResourceGuid(fmt.Sprintf("/v3/isolation_segments?names=%s", isolationSegmentName))
}

func attachIsolationSegmentToOrg(environment *ReproducibleTestSuiteSetup, isoSegGuid string) {
	orgGuid := getOrganizationGuid(environment.GetOrganizationName())
	response := cf.Cf(
		"curl",
		fmt.Sprintf("/v3/isolation_segments/%s/relationships/organizations", isoSegGuid),
		"-X",
		"POST",
		"-d",
		fmt.Sprintf(`{"data":[{"guid": "%s"}]}`, orgGuid),
	)
	Expect(response.Wait()).To(Exit(0))
}

func attachIsolationSegmentToSpace(environment *ReproducibleTestSuiteSetup, isoSegGuid string) {
	spaceGuid := getSpaceGuidForOrg(getOrganizationGuid(environment.GetOrganizationName()))
	response := cf.Cf(
		"curl",
		fmt.Sprintf("/v2/spaces/%s", spaceGuid),
		"-X",
		"PUT",
		"-d",
		fmt.Sprintf(`{"isolation_segment_guid":"%s"}`, isoSegGuid),
	)
	Expect(response.Wait()).To(Exit(0))
}

func getOrganizationGuid(organizationName string) string {
	return getV2ResourceGuid(fmt.Sprintf("/v2/organizations?q=name:%s", organizationName))
}

func getSpaceGuidForOrg(orgGuid string) string {
	return getV2ResourceGuid(fmt.Sprintf("/v2/organizations/%s/spaces", orgGuid))
}

func getV2ResourceGuid(endpoint string) string {
	response := cf.Cf("curl", endpoint, "-X", "GET")
	Expect(response.Wait()).To(Exit(0))
	var r struct {
		Resources []struct {
			Metadata struct {
				Guid string
			}
		}
	}
	Expect(json.Unmarshal(response.Out.Contents(), &r)).To(Succeed())
	Expect(len(r.Resources)).To(Equal(1))
	return r.Resources[0].Metadata.Guid
}

func getV3ResourceGuid(endpoint string) string {
	response := cf.Cf("curl", endpoint, "-X", "GET")
	Expect(response.Wait()).To(Exit(0))
	var r struct {
		Resources []struct {
			Guid string
		}
	}
	Expect(json.Unmarshal(response.Out.Contents(), &r)).To(Succeed())
	if len(r.Resources) == 0 {
		return ""
	}
	return r.Resources[0].Guid
}
<|endoftext|>


<commit_before>package metrics

import (
	"context"
	"log"
	"net/url"
	"sort"
	"sync"
	"time"

	"github.com/DataDog/datadog-api-client-go/api/v1/datadog"
	"github.com/DataDog/datadog-go/statsd"
)

// DatadogClient is a metrics.Client that emits directly to Datadog.
type DatadogClient struct {
	StatsdClient
	cfg *datadog.Configuration
}

var _ Client = (*DatadogClient)(nil)

func NewDatadogClient(apiKey, appKey string, opts ...DatadogClientOpt) *DatadogClient {
	cfg := datadog.NewConfiguration()
	c := &DatadogClient{
		cfg: cfg,
		StatsdClient: StatsdClient{
			client: newDatadogStatsd(cfg, apiKey, appKey),
		},
	}
	for _, o := range opts {
		o(c)
	}
	return c
}

type DatadogClientOpt func(*DatadogClient)

func WithDatadogTags(tags []string) DatadogClientOpt {
	return func(c *DatadogClient) {
		c.client.(*datadogStatsd).tags = tags
	}
}

func WithDatadogURL(datadog url.URL) DatadogClientOpt {
	return func(c *DatadogClient) {
		c.cfg.Host = datadog.Host
		c.cfg.Scheme = datadog.Scheme
	}
}

func WithDatadogFrozenClock(frozenTime float64) DatadogClientOpt {
	return func(c *DatadogClient) {
		c.client.(*datadogStatsd).now = func() float64 { return frozenTime }
	}
}

func (d *DatadogClient) Close() {
	d.client.(*datadogStatsd).submit()
}

// datadogStatsd is an alternative statsd.Client that transmits directly to Datadog
type datadogStatsd struct {
	authCtx context.Context
	metrics *datadog.MetricsApiService
	events  *datadog.EventsApiService
	tags    []string
	now     func() float64

	mu     sync.Mutex
	series []*datadog.Series
}

var _ statsdClient = (*datadogStatsd)(nil)

const submitTimeout = 3 * time.Second

func newDatadogStatsd(cfg *datadog.Configuration, apiKey, appKey string) *datadogStatsd {
	client := datadog.NewAPIClient(cfg)
	keys := map[string]datadog.APIKey{
		"apiKeyAuth": {Key: apiKey},
		"appKeyAuth": {Key: appKey},
	}
	return &datadogStatsd{
		authCtx: context.WithValue(context.Background(), datadog.ContextAPIKeys, keys),
		metrics: client.MetricsApi,
		events:  client.EventsApi,
		now:     func() float64 { return float64(time.Now().Unix()) },
	}
}

const (
	durationType = "gauge"
	countType    = "count"
)

func (d *datadogStatsd) Incr(metric string, tags []string, _ float64) error {
	tags = append(d.tags, tags...)
	now := d.now()

	d.mu.Lock()
	defer d.mu.Unlock()

	if existing := d.findSeries(metric, tags); existing != nil {
		for i, p := range existing.Points {
			if p[0] == now {
				existing.Points[i][1]++
				return nil
			}
		}
		existing.Points = append(existing.Points, []float64{now, 1})
		return nil
	}

	// Not found, create
	s := datadog.NewSeries(metric, [][]float64{{now, 1}})
	s.SetType(countType)
	s.SetTags(tags)
	d.series = append(d.series, s)
	return nil
}

func (d *datadogStatsd) Timing(metric string, dur time.Duration, tags []string, _ float64) error {
	tags = append(d.tags, tags...)
	now := d.now()
	val := float64(dur.Milliseconds())

	d.mu.Lock()
	defer d.mu.Unlock()

	if existing := d.findSeries(metric, tags); existing != nil {
		existing.Points = append(existing.Points, []float64{now, val})
		return nil
	}

	// Not found, create
	s := datadog.NewSeries(metric, [][]float64{{now, val}})
	s.SetType(durationType)
	s.SetTags(tags)
	d.series = append(d.series, s)
	return nil
}

func (d *datadogStatsd) Event(e *statsd.Event) error {
	ctx, cancel := context.WithTimeout(d.authCtx, submitTimeout)
	defer cancel()

	ddEvent := datadog.NewEventCreateRequest(e.Text, e.Title)
	ddEvent.SetAlertType(datadog.EventAlertType(e.AlertType))
	ddEvent.SetAggregationKey(e.AggregationKey)
	ddEvent.SetPriority(datadog.EventPriority(e.Priority))
	ddEvent.SetTags(append(d.tags, e.Tags...))
	if e.Timestamp.IsZero() {
		ddEvent.SetDateHappened(int64(d.now()))
	} else {
		ddEvent.SetDateHappened(e.Timestamp.Unix())
	}

	if _, _, err := d.events.CreateEvent(ctx, *ddEvent); err != nil {
		log.Println("error submitting event to datadog", err)
	}
	return nil
}

func (d *datadogStatsd) submit() {
	ctx, cancel := context.WithTimeout(d.authCtx, submitTimeout)
	defer cancel()

	d.mu.Lock()
	defer d.mu.Unlock()

	seriesCount := len(d.series)
	if seriesCount == 0 {
		return
	}
	series := make([]datadog.Series, 0, seriesCount)
	for _, s := range d.series {
		series = append(series, *s)
	}
	d.series = nil

	// TODO: this is sync, that is not great
	if _, _, err := d.metrics.SubmitMetrics(ctx, *datadog.NewMetricsPayload(series)); err != nil {
		log.Println("error submitting metrics to datadog", err)
	}
}

func (d *datadogStatsd) findSeries(metric string, tags []string) *datadog.Series {
	sort.Strings(tags)
seriesLoop:
	for _, s := range d.series {
		if s.GetMetric() != metric {
			continue
		}
		sTags := s.GetTags()
		if len(sTags) != len(tags) {
			continue
		}
		sort.Strings(sTags)
		for i, t := range sTags {
			if tags[i] != t {
				continue seriesLoop
			}
		}
		return s
	}
	return nil
}
<commit_msg>batched timings, mock dogstatsd aggregation<commit_after>package metrics

import (
	"context"
	"fmt"
	"log"
	"math"
	"net/url"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/DataDog/datadog-api-client-go/api/v1/datadog"
	"github.com/DataDog/datadog-go/statsd"
)

// DatadogClient is a metrics.Client that emits directly to Datadog.
type DatadogClient struct {
	StatsdClient
	cfg *datadog.Configuration
}

var _ Client = (*DatadogClient)(nil)

func NewDatadogClient(apiKey, appKey string, opts ...DatadogClientOpt) *DatadogClient {
	cfg := datadog.NewConfiguration()
	c := &DatadogClient{
		cfg: cfg,
		StatsdClient: StatsdClient{
			client: newDatadogStatsd(cfg, apiKey, appKey),
		},
	}
	for _, o := range opts {
		o(c)
	}
	return c
}

type DatadogClientOpt func(*DatadogClient)

func WithDatadogTags(tags []string) DatadogClientOpt {
	return func(c *DatadogClient) {
		c.client.(*datadogStatsd).tags = tags
	}
}

func WithDatadogURL(datadog url.URL) DatadogClientOpt {
	return func(c *DatadogClient) {
		c.cfg.Host = datadog.Host
		c.cfg.Scheme = datadog.Scheme
	}
}

func WithDatadogFrozenClock(frozenTime float64) DatadogClientOpt {
	return func(c *DatadogClient) {
		c.client.(*datadogStatsd).now = func() float64 { return frozenTime }
	}
}

func (d *DatadogClient) Close() {
	d.client.(*datadogStatsd).submit()
}

// datadogStatsd is an alternative statsd.Client that transmits directly to Datadog
type datadogStatsd struct {
	authCtx context.Context
	metrics *datadog.MetricsApiService
	events  *datadog.EventsApiService
	tags    []string
	now     func() float64

	mu     sync.Mutex
	series []*datadog.Series
	// map of timingKey() to timestamp, to values
	durationData map[string]map[float64][]float64
}

var _ statsdClient = (*datadogStatsd)(nil)

const submitTimeout = 3 * time.Second

func newDatadogStatsd(cfg *datadog.Configuration, apiKey, appKey string) *datadogStatsd {
	client := datadog.NewAPIClient(cfg)
	keys := map[string]datadog.APIKey{
		"apiKeyAuth": {Key: apiKey},
		"appKeyAuth": {Key: appKey},
	}
	return &datadogStatsd{
		authCtx:      context.WithValue(context.Background(), datadog.ContextAPIKeys, keys),
		metrics:      client.MetricsApi,
		events:       client.EventsApi,
		now:          func() float64 { return float64(time.Now().Unix()) },
		durationData: make(map[string]map[float64][]float64),
	}
}

const (
	durationType = "gauge"
	countType    = "count"
)

func (d *datadogStatsd) Incr(metric string, tags []string, _ float64) error {
	tags = append(d.tags, tags...)
	now := d.now()

	d.mu.Lock()
	defer d.mu.Unlock()

	if existing := d.findSeries(metric, tags); existing != nil {
		for i, p := range existing.Points {
			if p[0] == now {
				existing.Points[i][1]++
				return nil
			}
		}
		existing.Points = append(existing.Points, []float64{now, 1})
		return nil
	}

	// Not found, create
	s := datadog.NewSeries(metric, [][]float64{{now, 1}})
	s.SetType(countType)
	s.SetTags(tags)
	d.series = append(d.series, s)
	return nil
}

func (d *datadogStatsd) Timing(metric string, dur time.Duration, tags []string, _ float64) error {
	tags = append(d.tags, tags...)
	now := d.now()
	val := float64(dur.Milliseconds())
	tk := timingKey(metric, tags)

	d.mu.Lock()
	defer d.mu.Unlock()

	if existing := d.findSeries(metric, tags); existing != nil {
		d.durationData[tk][now] = append(d.durationData[tk][now], val)
		return nil
	}

	// Not found, create
	s := datadog.NewSeries(metric, nil)
	s.SetType(durationType)
	s.SetTags(tags)
	s.SetInterval(1)
	d.series = append(d.series, s)
	d.durationData[tk] = map[float64][]float64{now: {val}}
	return nil
}

func timingKey(metric string, tags []string) string {
	return fmt.Sprintf("%s %s", metric, strings.Join(tags, ","))
}

func (d *datadogStatsd) Event(e *statsd.Event) error {
	ctx, cancel := context.WithTimeout(d.authCtx, submitTimeout)
	defer cancel()

	ddEvent := datadog.NewEventCreateRequest(e.Text, e.Title)
	ddEvent.SetAlertType(datadog.EventAlertType(e.AlertType))
	ddEvent.SetAggregationKey(e.AggregationKey)
	ddEvent.SetPriority(datadog.EventPriority(e.Priority))
	ddEvent.SetTags(append(d.tags, e.Tags...))
	if e.Timestamp.IsZero() {
		ddEvent.SetDateHappened(int64(d.now()))
	} else {
		ddEvent.SetDateHappened(e.Timestamp.Unix())
	}

	if _, _, err := d.events.CreateEvent(ctx, *ddEvent); err != nil {
		log.Println("error submitting event to datadog", err)
	}
	return nil
}

func (d *datadogStatsd) submit() {
	series := d.flushSeries()
	if len(series) == 0 {
		return
	}

	ctx, cancel := context.WithTimeout(d.authCtx, submitTimeout)
	defer cancel()

	// TODO: this is sync, that is not great
	if _, _, err := d.metrics.SubmitMetrics(ctx, *datadog.NewMetricsPayload(series)); err != nil {
		log.Println("error submitting metrics to datadog", err)
	}
}

func (d *datadogStatsd) flushSeries() []datadog.Series {
	d.mu.Lock()
	defer d.mu.Unlock()

	series := make([]datadog.Series, 0, len(d.series))
	for _, s := range d.series {
		if s.GetType() != durationType {
			series = append(series, *s)
			continue
		}

		// Construct series from captured samples:
		tk := timingKey(s.GetMetric(), s.GetTags())
		var counts, averages, p95s, p99s, maxes [][]float64
		for ts, data := range d.durationData[tk] {
			sort.Float64s(data)
			var sum float64
			for _, f := range data {
				sum += f
			}
			count := float64(len(data))
			counts = append(counts, []float64{ts, count})
			averages = append(averages, []float64{ts, sum / count})
			p95s = append(p95s, []float64{ts, percentile(data, 0.95)})
			p99s = append(p99s, []float64{ts, percentile(data, 0.99)})
			maxes = append(maxes, []float64{ts, data[len(data)-1]})
		}
		series = append(series, cloneSeries(s, "count", counts))
		series = append(series, cloneSeries(s, "avg", averages))
		series = append(series, cloneSeries(s, "95percentile", p95s))
		series = append(series, cloneSeries(s, "99percentile", p99s))
		series = append(series, cloneSeries(s, "max", maxes))
	}
	d.series = nil
	d.durationData = make(map[string]map[float64][]float64)
	return series
}

func cloneSeries(s *datadog.Series, suffix string, points [][]float64) datadog.Series {
	return datadog.Series{
		Host:           s.Host,
		Interval:       s.Interval,
		Metric:         fmt.Sprintf("%s.%s", s.GetMetric(), suffix),
		Points:         points,
		Tags:           s.Tags,
		Type:           s.Type,
		UnparsedObject: s.UnparsedObject,
	}
}

func percentile(data []float64, p float64) float64 {
	pos := float64(len(data)) * p
	if math.Round(pos) == pos {
		// Return exact value at percentile
		return data[int(pos)]
	}

	return (data[int(pos-1)] + data[int(pos)]) / 2
}

func (d *datadogStatsd) findSeries(metric string, tags []string) *datadog.Series {
	sort.Strings(tags)
seriesLoop:
	for _, s := range d.series {
		if s.GetMetric() != metric {
			continue
		}
		sTags := s.GetTags()
		if len(sTags) != len(tags) {
			continue
		}
		sort.Strings(sTags)
		for i, t := range sTags {
			if tags[i] != t {
				continue seriesLoop
			}
		}
		return s
	}
	return nil
}
<|endoftext|>