<commit_before>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cmd

import (
	"bytes"
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/log"
	"sort"
	"strings"
)

type topic struct {
	short string
	long  func() string
}

// MissingCallback defines a function that will be used by the SuperCommand if
// the requested subcommand isn't found.
type MissingCallback func(ctx *Context, subcommand string, args []string) error

// SuperCommandParams provides a way to have default parameter to the
// `NewSuperCommand` call.
type SuperCommandParams struct {
	Name            string
	Purpose         string
	Doc             string
	Log             *Log
	MissingCallback MissingCallback
}

// NewSuperCommand creates and initializes a new `SuperCommand`, and returns
// the fully initialized structure.
func NewSuperCommand(params SuperCommandParams) *SuperCommand {
	command := &SuperCommand{
		Name:            params.Name,
		Purpose:         params.Purpose,
		Doc:             params.Doc,
		Log:             params.Log,
		missingCallback: params.MissingCallback}
	command.init()
	return command
}

// SuperCommand is a Command that selects a subcommand and assumes its
// properties; any command line arguments that were not used in selecting
// the subcommand are passed down to it, and to Run a SuperCommand is to run
// its selected subcommand.
type SuperCommand struct {
	CommandBase
	Name            string
	Purpose         string
	Doc             string
	Log             *Log
	subcmds         map[string]Command
	flags           *gnuflag.FlagSet
	subcmd          Command
	showHelp        bool
	missingCallback MissingCallback
}

// Because Go doesn't have constructors that initialize the object into a
// ready state.
func (c *SuperCommand) init() {
	if c.subcmds != nil {
		return
	}
	help := &helpCommand{
		super: c,
	}
	help.init()
	c.subcmds = map[string]Command{
		"help": help,
	}
}

func (c *SuperCommand) AddHelpTopic(name, short, long string) {
	c.subcmds["help"].(*helpCommand).addTopic(name, short, long)
}

// Register makes a subcommand available for use on the command line. The
// command will be available via its own name, and via any supplied aliases.
func (c *SuperCommand) Register(subcmd Command) {
	info := subcmd.Info()
	c.insert(info.Name, subcmd)
	for _, name := range info.Aliases {
		c.insert(name, subcmd)
	}
}

func (c *SuperCommand) insert(name string, subcmd Command) {
	if _, found := c.subcmds[name]; found || name == "help" {
		panic(fmt.Sprintf("command already registered: %s", name))
	}
	c.subcmds[name] = subcmd
}

// describeCommands returns a short description of each registered subcommand.
func (c *SuperCommand) describeCommands(simple bool) string {
	var lineFormat = "    %-*s - %s"
	var outputFormat = "commands:\n%s"
	if simple {
		lineFormat = "%-*s  %s"
		outputFormat = "%s"
	}
	cmds := make([]string, len(c.subcmds))
	i := 0
	longest := 0
	for name := range c.subcmds {
		if len(name) > longest {
			longest = len(name)
		}
		cmds[i] = name
		i++
	}
	sort.Strings(cmds)
	for i, name := range cmds {
		info := c.subcmds[name].Info()
		purpose := info.Purpose
		if name != info.Name {
			purpose = "alias for " + info.Name
		}
		cmds[i] = fmt.Sprintf(lineFormat, longest, name, purpose)
	}
	return fmt.Sprintf(outputFormat, strings.Join(cmds, "\n"))
}

// Info returns a description of the currently selected subcommand, or of the
// SuperCommand itself if no subcommand has been specified.
func (c *SuperCommand) Info() *Info {
	if c.subcmd != nil {
		info := *c.subcmd.Info()
		info.Name = fmt.Sprintf("%s %s", c.Name, info.Name)
		return &info
	}
	docParts := []string{}
	if doc := strings.TrimSpace(c.Doc); doc != "" {
		docParts = append(docParts, doc)
	}
	if cmds := c.describeCommands(false); cmds != "" {
		docParts = append(docParts, cmds)
	}
	return &Info{
		Name:    c.Name,
		Args:    "<command> ...",
		Purpose: c.Purpose,
		Doc:     strings.Join(docParts, "\n\n"),
	}
}

const helpPurpose = "show help on a command or other topic"

// SetFlags adds the options that apply to all commands, particularly those
// due to logging.
func (c *SuperCommand) SetFlags(f *gnuflag.FlagSet) {
	if c.Log != nil {
		c.Log.AddFlags(f)
	}
	f.BoolVar(&c.showHelp, "h", false, helpPurpose)
	f.BoolVar(&c.showHelp, "help", false, "")

	c.flags = f
}

// Init initializes the command for running.
func (c *SuperCommand) Init(args []string) error {
	if len(args) == 0 {
		c.subcmd = c.subcmds["help"]
		return nil
	}

	found := false
	// Look for the command.
	if c.subcmd, found = c.subcmds[args[0]]; !found {
		if c.missingCallback != nil {
			c.subcmd = &missingCommand{
				callback: c.missingCallback,
				name:     args[0],
				args:     args[1:],
			}
			// Yes return here, no Init called on missing Command.
			return nil
		}
		return fmt.Errorf("unrecognized command: %s %s", c.Info().Name, args[0])
	}
	args = args[1:]
	c.subcmd.SetFlags(c.flags)
	if err := c.flags.Parse(true, args); err != nil {
		return err
	}
	args = c.flags.Args()
	if c.showHelp {
		// We want to treat help for the command the same way we would if we went "help foo".
		args = []string{c.subcmd.Info().Name}
		c.subcmd = c.subcmds["help"]
	}
	return c.subcmd.Init(args)
}

// Run executes the subcommand that was selected in Init.
func (c *SuperCommand) Run(ctx *Context) error {
	if c.subcmd == nil {
		panic("Run: missing subcommand; Init failed or not called")
	}
	if c.Log != nil {
		if err := c.Log.Start(ctx); err != nil {
			return err
		}
	}
	err := c.subcmd.Run(ctx)
	if err != nil && err != ErrSilent {
		log.Errorf("command failed: %v", err)
	} else {
		log.Infof("command finished")
	}
	return err
}

type missingCommand struct {
	CommandBase
	callback MissingCallback
	name     string
	args     []string
}

// Missing commands only need to supply Init for the interface, but this is
// never called.
func (c *missingCommand) Info() *Info {
	return nil
}

func (c *missingCommand) Run(ctx *Context) error {
	return c.callback(ctx, c.name, c.args)
}

type helpCommand struct {
	CommandBase
	super  *SuperCommand
	topic  string
	topics map[string]topic
}

func (c *helpCommand) init() {
	c.topics = map[string]topic{
		"commands": {
			short: "Basic help for all commands",
			long:  func() string { return c.super.describeCommands(true) },
		},
		"global-options": {
			short: "Options common to all commands",
			long:  func() string { return c.globalOptions() },
		},
		"topics": {
			short: "Topic list",
			long:  func() string { return c.topicList() },
		},
	}
}

func echo(s string) func() string {
	return func() string { return s }
}

func (c *helpCommand) addTopic(name, short, long string) {
	if _, found := c.topics[name]; found {
		panic(fmt.Sprintf("help topic already added: %s", name))
	}
	c.topics[name] = topic{short, echo(long)}
}

func (c *helpCommand) globalOptions() string {
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, `Global Options

These options may be used with any command, and may appear in front of any
command.

`)

	f := gnuflag.NewFlagSet("", gnuflag.ContinueOnError)
	c.super.SetFlags(f)
	f.SetOutput(buf)
	f.PrintDefaults()
	return buf.String()
}

func (c *helpCommand) topicList() string {
	topics := make([]string, len(c.topics))
	i := 0
	longest := 0
	for name := range c.topics {
		if len(name) > longest {
			longest = len(name)
		}
		topics[i] = name
		i++
	}
	sort.Strings(topics)
	for i, name := range topics {
		shortHelp := c.topics[name].short
		topics[i] = fmt.Sprintf("%-*s  %s", longest, name, shortHelp)
	}
	return fmt.Sprintf("%s", strings.Join(topics, "\n"))
}

func (c *helpCommand) Info() *Info {
	return &Info{
		Name:    "help",
		Args:    "[topic]",
		Purpose: helpPurpose,
		Doc: `
See also: topics
`,
	}
}

func (c *helpCommand) Init(args []string) error {
	switch len(args) {
	case 0:
	case 1:
		c.topic = args[0]
	default:
		return fmt.Errorf("extra arguments to command help: %q", args[2:])
	}
	return nil
}

func (c *helpCommand) Run(ctx *Context) error {
	// If there is no help topic specified, print basic usage.
	if c.topic == "" {
		if _, ok := c.topics["basics"]; ok {
			c.topic = "basics"
		} else {
			// At this point, "help" is selected as the SuperCommand's
			// sub-command, but we want the info to be printed
			// as if there was nothing selected.
			c.super.subcmd = nil

			info := c.super.Info()
			f := gnuflag.NewFlagSet(info.Name, gnuflag.ContinueOnError)
			c.SetFlags(f)
			ctx.Stdout.Write(info.Help(f))
			return nil
		}
	}
	if helpcmd, ok := c.super.subcmds[c.topic]; ok {
		info := helpcmd.Info()
		info.Name = fmt.Sprintf("%s %s", c.super.Name, info.Name)
		f := gnuflag.NewFlagSet(info.Name, gnuflag.ContinueOnError)
		helpcmd.SetFlags(f)
		ctx.Stdout.Write(info.Help(f))
		return nil
	}
	topic, ok := c.topics[c.topic]
	if !ok {
		return fmt.Errorf("unknown command or topic for %s", c.topic)
	}
	fmt.Fprintf(ctx.Stdout, "%s\n", strings.TrimSpace(topic.long()))
	return nil
}
<commit_msg>Fix a typo.<commit_after>// Copyright 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cmd

import (
	"bytes"
	"fmt"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/log"
	"sort"
	"strings"
)

type topic struct {
	short string
	long  func() string
}

// MissingCallback defines a function that will be used by the SuperCommand if
// the requested subcommand isn't found.
type MissingCallback func(ctx *Context, subcommand string, args []string) error

// SuperCommandParams provides a way to have default parameter to the
// `NewSuperCommand` call.
type SuperCommandParams struct {
	Name            string
	Purpose         string
	Doc             string
	Log             *Log
	MissingCallback MissingCallback
}

// NewSuperCommand creates and initializes a new `SuperCommand`, and returns
// the fully initialized structure.
func NewSuperCommand(params SuperCommandParams) *SuperCommand {
	command := &SuperCommand{
		Name:            params.Name,
		Purpose:         params.Purpose,
		Doc:             params.Doc,
		Log:             params.Log,
		missingCallback: params.MissingCallback}
	command.init()
	return command
}

// SuperCommand is a Command that selects a subcommand and assumes its
// properties; any command line arguments that were not used in selecting
// the subcommand are passed down to it, and to Run a SuperCommand is to run
// its selected subcommand.
type SuperCommand struct {
	CommandBase
	Name            string
	Purpose         string
	Doc             string
	Log             *Log
	subcmds         map[string]Command
	flags           *gnuflag.FlagSet
	subcmd          Command
	showHelp        bool
	missingCallback MissingCallback
}

// Because Go doesn't have constructors that initialize the object into a
// ready state.
func (c *SuperCommand) init() {
	if c.subcmds != nil {
		return
	}
	help := &helpCommand{
		super: c,
	}
	help.init()
	c.subcmds = map[string]Command{
		"help": help,
	}
}

func (c *SuperCommand) AddHelpTopic(name, short, long string) {
	c.subcmds["help"].(*helpCommand).addTopic(name, short, long)
}

// Register makes a subcommand available for use on the command line. The
// command will be available via its own name, and via any supplied aliases.
func (c *SuperCommand) Register(subcmd Command) {
	info := subcmd.Info()
	c.insert(info.Name, subcmd)
	for _, name := range info.Aliases {
		c.insert(name, subcmd)
	}
}

func (c *SuperCommand) insert(name string, subcmd Command) {
	if _, found := c.subcmds[name]; found || name == "help" {
		panic(fmt.Sprintf("command already registered: %s", name))
	}
	c.subcmds[name] = subcmd
}

// describeCommands returns a short description of each registered subcommand.
func (c *SuperCommand) describeCommands(simple bool) string {
	var lineFormat = "    %-*s - %s"
	var outputFormat = "commands:\n%s"
	if simple {
		lineFormat = "%-*s  %s"
		outputFormat = "%s"
	}
	cmds := make([]string, len(c.subcmds))
	i := 0
	longest := 0
	for name := range c.subcmds {
		if len(name) > longest {
			longest = len(name)
		}
		cmds[i] = name
		i++
	}
	sort.Strings(cmds)
	for i, name := range cmds {
		info := c.subcmds[name].Info()
		purpose := info.Purpose
		if name != info.Name {
			purpose = "alias for " + info.Name
		}
		cmds[i] = fmt.Sprintf(lineFormat, longest, name, purpose)
	}
	return fmt.Sprintf(outputFormat, strings.Join(cmds, "\n"))
}

// Info returns a description of the currently selected subcommand, or of the
// SuperCommand itself if no subcommand has been specified.
func (c *SuperCommand) Info() *Info {
	if c.subcmd != nil {
		info := *c.subcmd.Info()
		info.Name = fmt.Sprintf("%s %s", c.Name, info.Name)
		return &info
	}
	docParts := []string{}
	if doc := strings.TrimSpace(c.Doc); doc != "" {
		docParts = append(docParts, doc)
	}
	if cmds := c.describeCommands(false); cmds != "" {
		docParts = append(docParts, cmds)
	}
	return &Info{
		Name:    c.Name,
		Args:    "<command> ...",
		Purpose: c.Purpose,
		Doc:     strings.Join(docParts, "\n\n"),
	}
}

const helpPurpose = "show help on a command or other topic"

// SetFlags adds the options that apply to all commands, particularly those
// due to logging.
func (c *SuperCommand) SetFlags(f *gnuflag.FlagSet) {
	if c.Log != nil {
		c.Log.AddFlags(f)
	}
	f.BoolVar(&c.showHelp, "h", false, helpPurpose)
	f.BoolVar(&c.showHelp, "help", false, "")

	c.flags = f
}

// Init initializes the command for running.
func (c *SuperCommand) Init(args []string) error {
	if len(args) == 0 {
		c.subcmd = c.subcmds["help"]
		return nil
	}

	found := false
	// Look for the command.
	if c.subcmd, found = c.subcmds[args[0]]; !found {
		if c.missingCallback != nil {
			c.subcmd = &missingCommand{
				callback: c.missingCallback,
				name:     args[0],
				args:     args[1:],
			}
			// Yes return here, no Init called on missing Command.
			return nil
		}
		return fmt.Errorf("unrecognized command: %s %s", c.Info().Name, args[0])
	}
	args = args[1:]
	c.subcmd.SetFlags(c.flags)
	if err := c.flags.Parse(true, args); err != nil {
		return err
	}
	args = c.flags.Args()
	if c.showHelp {
		// We want to treat help for the command the same way we would if we went "help foo".
		args = []string{c.subcmd.Info().Name}
		c.subcmd = c.subcmds["help"]
	}
	return c.subcmd.Init(args)
}

// Run executes the subcommand that was selected in Init.
func (c *SuperCommand) Run(ctx *Context) error {
	if c.subcmd == nil {
		panic("Run: missing subcommand; Init failed or not called")
	}
	if c.Log != nil {
		if err := c.Log.Start(ctx); err != nil {
			return err
		}
	}
	err := c.subcmd.Run(ctx)
	if err != nil && err != ErrSilent {
		log.Errorf("command failed: %v", err)
	} else {
		log.Infof("command finished")
	}
	return err
}

type missingCommand struct {
	CommandBase
	callback MissingCallback
	name     string
	args     []string
}

// Missing commands only need to supply Info for the interface, but this is
// never called.
func (c *missingCommand) Info() *Info {
	return nil
}

func (c *missingCommand) Run(ctx *Context) error {
	return c.callback(ctx, c.name, c.args)
}

type helpCommand struct {
	CommandBase
	super  *SuperCommand
	topic  string
	topics map[string]topic
}

func (c *helpCommand) init() {
	c.topics = map[string]topic{
		"commands": {
			short: "Basic help for all commands",
			long:  func() string { return c.super.describeCommands(true) },
		},
		"global-options": {
			short: "Options common to all commands",
			long:  func() string { return c.globalOptions() },
		},
		"topics": {
			short: "Topic list",
			long:  func() string { return c.topicList() },
		},
	}
}

func echo(s string) func() string {
	return func() string { return s }
}

func (c *helpCommand) addTopic(name, short, long string) {
	if _, found := c.topics[name]; found {
		panic(fmt.Sprintf("help topic already added: %s", name))
	}
	c.topics[name] = topic{short, echo(long)}
}

func (c *helpCommand) globalOptions() string {
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, `Global Options

These options may be used with any command, and may appear in front of any
command.

`)

	f := gnuflag.NewFlagSet("", gnuflag.ContinueOnError)
	c.super.SetFlags(f)
	f.SetOutput(buf)
	f.PrintDefaults()
	return buf.String()
}

func (c *helpCommand) topicList() string {
	topics := make([]string, len(c.topics))
	i := 0
	longest := 0
	for name := range c.topics {
		if len(name) > longest {
			longest = len(name)
		}
		topics[i] = name
		i++
	}
	sort.Strings(topics)
	for i, name := range topics {
		shortHelp := c.topics[name].short
		topics[i] = fmt.Sprintf("%-*s  %s", longest, name, shortHelp)
	}
	return fmt.Sprintf("%s", strings.Join(topics, "\n"))
}

func (c *helpCommand) Info() *Info {
	return &Info{
		Name:    "help",
		Args:    "[topic]",
		Purpose: helpPurpose,
		Doc: `
See also: topics
`,
	}
}

func (c *helpCommand) Init(args []string) error {
	switch len(args) {
	case 0:
	case 1:
		c.topic = args[0]
	default:
		return fmt.Errorf("extra arguments to command help: %q", args[2:])
	}
	return nil
}

func (c *helpCommand) Run(ctx *Context) error {
	// If there is no help topic specified, print basic usage.
	if c.topic == "" {
		if _, ok := c.topics["basics"]; ok {
			c.topic = "basics"
		} else {
			// At this point, "help" is selected as the SuperCommand's
			// sub-command, but we want the info to be printed
			// as if there was nothing selected.
			c.super.subcmd = nil

			info := c.super.Info()
			f := gnuflag.NewFlagSet(info.Name, gnuflag.ContinueOnError)
			c.SetFlags(f)
			ctx.Stdout.Write(info.Help(f))
			return nil
		}
	}
	if helpcmd, ok := c.super.subcmds[c.topic]; ok {
		info := helpcmd.Info()
		info.Name = fmt.Sprintf("%s %s", c.super.Name, info.Name)
		f := gnuflag.NewFlagSet(info.Name, gnuflag.ContinueOnError)
		helpcmd.SetFlags(f)
		ctx.Stdout.Write(info.Help(f))
		return nil
	}
	topic, ok := c.topics[c.topic]
	if !ok {
		return fmt.Errorf("unknown command or topic for %s", c.topic)
	}
	fmt.Fprintf(ctx.Stdout, "%s\n", strings.TrimSpace(topic.long()))
	return nil
}
<|endoftext|>


<commit_before>// Downloads torrents from the command-line.
package main

import (
	"expvar"
	"fmt"
	"io"
	stdLog "log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/alexflint/go-arg"
	"github.com/anacrolix/missinggo"
	"github.com/anacrolix/torrent/bencode"
	"github.com/davecgh/go-spew/spew"
	"github.com/dustin/go-humanize"
	"golang.org/x/xerrors"

	"github.com/anacrolix/log"

	"github.com/anacrolix/envpprof"
	"github.com/anacrolix/tagflag"
	"golang.org/x/time/rate"

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/iplist"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
)

func torrentBar(t *torrent.Torrent, pieceStates bool) {
	go func() {
		start := time.Now()
		if t.Info() == nil {
			fmt.Printf("%v: getting torrent info for %q\n", time.Since(start), t.Name())
			<-t.GotInfo()
		}
		var lastLine string
		for range time.Tick(time.Second) {
			var completedPieces, partialPieces int
			psrs := t.PieceStateRuns()
			for _, r := range psrs {
				if r.Complete {
					completedPieces += r.Length
				}
				if r.Partial {
					partialPieces += r.Length
				}
			}
			line := fmt.Sprintf(
				"%v: downloading %q: %s/%s, %d/%d pieces completed (%d partial)\n",
				time.Since(start),
				t.Name(),
				humanize.Bytes(uint64(t.BytesCompleted())),
				humanize.Bytes(uint64(t.Length())),
				completedPieces,
				t.NumPieces(),
				partialPieces,
			)
			if line != lastLine {
				lastLine = line
				os.Stdout.WriteString(line)
			}
			if pieceStates {
				fmt.Println(psrs)
			}
		}
	}()
}

type stringAddr string

func (stringAddr) Network() string   { return "" }
func (me stringAddr) String() string { return string(me) }

func resolveTestPeers(addrs []string) (ret []torrent.PeerInfo) {
	for _, ta := range flags.TestPeer {
		ret = append(ret, torrent.PeerInfo{
			Addr: stringAddr(ta),
		})
	}
	return
}

func addTorrents(client *torrent.Client) error {
	testPeers := resolveTestPeers(flags.TestPeer)
	for _, arg := range flags.Torrent {
		t, err := func() (*torrent.Torrent, error) {
			if strings.HasPrefix(arg, "magnet:") {
				t, err := client.AddMagnet(arg)
				if err != nil {
					return nil, xerrors.Errorf("error adding magnet: %w", err)
				}
				return t, nil
			} else if strings.HasPrefix(arg, "http://") || strings.HasPrefix(arg, "https://") {
				response, err := http.Get(arg)
				if err != nil {
					return nil, xerrors.Errorf("Error downloading torrent file: %s", err)
				}

				metaInfo, err := metainfo.Load(response.Body)
				defer response.Body.Close()
				if err != nil {
					return nil, xerrors.Errorf("error loading torrent file %q: %s\n", arg, err)
				}
				t, err := client.AddTorrent(metaInfo)
				if err != nil {
					return nil, xerrors.Errorf("adding torrent: %w", err)
				}
				return t, nil
			} else if strings.HasPrefix(arg, "infohash:") {
				t, _ := client.AddTorrentInfoHash(metainfo.NewHashFromHex(strings.TrimPrefix(arg, "infohash:")))
				return t, nil
			} else {
				metaInfo, err := metainfo.LoadFromFile(arg)
				if err != nil {
					return nil, xerrors.Errorf("error loading torrent file %q: %s\n", arg, err)
				}
				t, err := client.AddTorrent(metaInfo)
				if err != nil {
					return nil, xerrors.Errorf("adding torrent: %w", err)
				}
				return t, nil
			}
		}()
		if err != nil {
			return xerrors.Errorf("adding torrent for %q: %w", arg, err)
		}
		if flags.Progress {
			torrentBar(t, flags.PieceStates)
		}
		t.AddPeers(testPeers)
		go func() {
			<-t.GotInfo()
			if len(flags.File) == 0 {
				t.DownloadAll()
			} else {
				for _, f := range t.Files() {
					for _, fileArg := range flags.File {
						if f.DisplayPath() == fileArg {
							f.Download()
						}
					}
				}
			}
		}()
	}
	return nil
}

var flags struct {
	Debug bool
	Stats *bool

	*DownloadCmd      `arg:"subcommand:download"`
	*ListFilesCmd     `arg:"subcommand:list-files"`
	*SpewBencodingCmd `arg:"subcommand:spew-bencoding"`
	*AnnounceCmd      `arg:"subcommand:announce"`
}

type SpewBencodingCmd struct{}

type DownloadCmd struct {
	Mmap            bool           `help:"memory-map torrent data"`
	TestPeer        []string       `help:"addresses of some starting peers"`
	Seed            bool           `help:"seed after download is complete"`
	Addr            string         `help:"network listen addr"`
	UploadRate      *tagflag.Bytes `help:"max piece bytes to send per second"`
	DownloadRate    *tagflag.Bytes `help:"max bytes per second down from peers"`
	PackedBlocklist string
	PublicIP        net.IP
	Progress        bool `default:"true"`
	PieceStates     bool
	Quiet           bool `help:"discard client logging"`
	Dht             bool `default:"true"`

	TcpPeers        bool `default:"true"`
	UtpPeers        bool `default:"true"`
	Webtorrent      bool `default:"true"`
	DisableWebseeds bool

	Ipv4 bool `default:"true"`
	Ipv6 bool `default:"true"`
	Pex  bool `default:"true"`

	File    []string
	Torrent []string `arity:"+" help:"torrent file path or magnet uri" arg:"positional"`
}

type ListFilesCmd struct {
	TorrentPath string `arg:"positional"`
}

func stdoutAndStderrAreSameFile() bool {
	fi1, _ := os.Stdout.Stat()
	fi2, _ := os.Stderr.Stat()
	return os.SameFile(fi1, fi2)
}

func statsEnabled() bool {
	if flags.Stats == nil {
		return flags.Debug
	}
	return *flags.Stats
}

func exitSignalHandlers(notify *missinggo.SynchronizedEvent) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	for {
		log.Printf("close signal received: %+v", <-c)
		notify.Set()
	}
}

func main() {
	if err := mainErr(); err != nil {
		log.Printf("error in main: %v", err)
		os.Exit(1)
	}
}

func mainErr() error {
	stdLog.SetFlags(stdLog.Flags() | stdLog.Lshortfile)
	p := arg.MustParse(&flags)
	switch {
	case flags.AnnounceCmd != nil:
		return announceErr()
	//case :
	//	return announceErr(flags.Args, parser)
	case flags.DownloadCmd != nil:
		return downloadErr()
	case flags.ListFilesCmd != nil:
		mi, err := metainfo.LoadFromFile(flags.ListFilesCmd.TorrentPath)
		if err != nil {
			return fmt.Errorf("loading from file %q: %v", flags.ListFilesCmd.TorrentPath, err)
		}
		info, err := mi.UnmarshalInfo()
		if err != nil {
			return fmt.Errorf("unmarshalling info from metainfo at %q: %v", flags.ListFilesCmd.TorrentPath, err)
		}
		for _, f := range info.UpvertedFiles() {
			fmt.Println(f.DisplayPath(&info))
		}
		return nil
	case flags.SpewBencodingCmd != nil:
		d := bencode.NewDecoder(os.Stdin)
		for i := 0; ; i++ {
			var v interface{}
			err := d.Decode(&v)
			if err == io.EOF {
				break
			}
			if err != nil {
				return fmt.Errorf("decoding message index %d: %w", i, err)
			}
			spew.Dump(v)
		}
		return nil
	default:
		p.Fail(fmt.Sprintf("unexpected subcommand: %v", p.Subcommand()))
		panic("unreachable")
	}
}

func downloadErr() error {
	defer envpprof.Stop()
	clientConfig := torrent.NewDefaultClientConfig()
	clientConfig.DisableWebseeds = flags.DisableWebseeds
	clientConfig.DisableTCP = !flags.TcpPeers
	clientConfig.DisableUTP = !flags.UtpPeers
	clientConfig.DisableIPv4 = !flags.Ipv4
	clientConfig.DisableIPv6 = !flags.Ipv6
	clientConfig.DisableAcceptRateLimiting = true
	clientConfig.NoDHT = !flags.Dht
	clientConfig.Debug = flags.Debug
	clientConfig.Seed = flags.Seed
	clientConfig.PublicIp4 = flags.PublicIP
	clientConfig.PublicIp6 = flags.PublicIP
	clientConfig.DisablePEX = !flags.Pex
	clientConfig.DisableWebtorrent = !flags.Webtorrent
	if flags.PackedBlocklist != "" {
		blocklist, err := iplist.MMapPackedFile(flags.PackedBlocklist)
		if err != nil {
			return xerrors.Errorf("loading blocklist: %v", err)
		}
		defer blocklist.Close()
		clientConfig.IPBlocklist = blocklist
	}
	if flags.Mmap {
		clientConfig.DefaultStorage = storage.NewMMap("")
	}
	if flags.Addr != "" {
		clientConfig.SetListenAddr(flags.Addr)
	}
	if flags.UploadRate != nil {
		clientConfig.UploadRateLimiter = rate.NewLimiter(rate.Limit(*flags.UploadRate), 256<<10)
	}
	if flags.DownloadRate != nil {
		clientConfig.DownloadRateLimiter = rate.NewLimiter(rate.Limit(*flags.DownloadRate), 1<<20)
	}
	if flags.Quiet {
		clientConfig.Logger = log.Discard
	}

	var stop missinggo.SynchronizedEvent
	defer func() {
		stop.Set()
	}()

	client, err := torrent.NewClient(clientConfig)
	if err != nil {
		return xerrors.Errorf("creating client: %v", err)
	}
	defer client.Close()
	go exitSignalHandlers(&stop)
	go func() {
		<-stop.C()
		client.Close()
	}()

	// Write status on the root path on the default HTTP muxer. This will be bound to localhost
	// somewhere if GOPPROF is set, thanks to the envpprof import.
	http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		client.WriteStatus(w)
	})
	err = addTorrents(client)
	if err != nil {
		return fmt.Errorf("adding torrents: %w", err)
	}
	defer outputStats(client)
	if client.WaitAll() {
		log.Print("downloaded ALL the torrents")
	} else {
		return xerrors.New("y u no complete torrents?!")
	}
	if flags.Seed {
		outputStats(client)
		<-stop.C()
	}
	return nil
}

func outputStats(cl *torrent.Client) {
	if !statsEnabled() {
		return
	}
	expvar.Do(func(kv expvar.KeyValue) {
		fmt.Printf("%s: %s\n", kv.Key, kv.Value)
	})
	cl.WriteStatus(os.Stdout)
}
<commit_msg>cmd/torrent: Track download rate<commit_after>// Downloads torrents from the command-line.
package main

import (
	"expvar"
	"fmt"
	"io"
	stdLog "log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/alexflint/go-arg"
	"github.com/anacrolix/missinggo"
	"github.com/anacrolix/torrent/bencode"
	"github.com/davecgh/go-spew/spew"
	"github.com/dustin/go-humanize"
	"golang.org/x/xerrors"

	"github.com/anacrolix/log"

	"github.com/anacrolix/envpprof"
	"github.com/anacrolix/tagflag"
	"golang.org/x/time/rate"

	"github.com/anacrolix/torrent"
	"github.com/anacrolix/torrent/iplist"
	"github.com/anacrolix/torrent/metainfo"
	"github.com/anacrolix/torrent/storage"
)

func torrentBar(t *torrent.Torrent, pieceStates bool) {
	go func() {
		start := time.Now()
		if t.Info() == nil {
			fmt.Printf("%v: getting torrent info for %q\n", time.Since(start), t.Name())
			<-t.GotInfo()
		}
		lastStats := t.Stats()
		var lastLine string
		for range time.Tick(time.Second) {
			var completedPieces, partialPieces int
			psrs := t.PieceStateRuns()
			for _, r := range psrs {
				if r.Complete {
					completedPieces += r.Length
				}
				if r.Partial {
					partialPieces += r.Length
				}
			}
			stats := t.Stats()
			line := fmt.Sprintf(
				"%v: downloading %q: %s/%s, %d/%d pieces completed (%d partial): %v/s\n",
				time.Since(start),
				t.Name(),
				humanize.Bytes(uint64(t.BytesCompleted())),
				humanize.Bytes(uint64(t.Length())),
				completedPieces,
				t.NumPieces(),
				partialPieces,
				humanize.Bytes(uint64(stats.BytesReadUsefulData.Int64()-lastStats.BytesReadUsefulData.Int64())),
			)
			if line != lastLine {
				lastLine = line
				os.Stdout.WriteString(line)
			}
			if pieceStates {
				fmt.Println(psrs)
			}
			lastStats = stats
		}
	}()
}

type stringAddr string

func (stringAddr) Network() string   { return "" }
func (me stringAddr) String() string { return string(me) }

func resolveTestPeers(addrs []string) (ret []torrent.PeerInfo) {
	for _, ta := range flags.TestPeer {
		ret = append(ret, torrent.PeerInfo{
			Addr: stringAddr(ta),
		})
	}
	return
}

func addTorrents(client *torrent.Client) error {
	testPeers := resolveTestPeers(flags.TestPeer)
	for _, arg := range flags.Torrent {
		t, err := func() (*torrent.Torrent, error) {
			if strings.HasPrefix(arg, "magnet:") {
				t, err := client.AddMagnet(arg)
				if err != nil {
					return nil, xerrors.Errorf("error adding magnet: %w", err)
				}
				return t, nil
			} else if strings.HasPrefix(arg, "http://") || strings.HasPrefix(arg, "https://") {
				response, err := http.Get(arg)
				if err != nil {
					return nil, xerrors.Errorf("Error downloading torrent file: %s", err)
				}

				metaInfo, err := metainfo.Load(response.Body)
				defer response.Body.Close()
				if err != nil {
					return nil, xerrors.Errorf("error loading torrent file %q: %s\n", arg, err)
				}
				t, err := client.AddTorrent(metaInfo)
				if err != nil {
					return nil, xerrors.Errorf("adding torrent: %w", err)
				}
				return t, nil
			} else if strings.HasPrefix(arg, "infohash:") {
				t, _ := client.AddTorrentInfoHash(metainfo.NewHashFromHex(strings.TrimPrefix(arg, "infohash:")))
				return t, nil
			} else {
				metaInfo, err := metainfo.LoadFromFile(arg)
				if err != nil {
					return nil, xerrors.Errorf("error loading torrent file %q: %s\n", arg, err)
				}
				t, err := client.AddTorrent(metaInfo)
				if err != nil {
					return nil, xerrors.Errorf("adding torrent: %w", err)
				}
				return t, nil
			}
		}()
		if err != nil {
			return xerrors.Errorf("adding torrent for %q: %w", arg, err)
		}
		if flags.Progress {
			torrentBar(t, flags.PieceStates)
		}
		t.AddPeers(testPeers)
		go func() {
			<-t.GotInfo()
			if len(flags.File) == 0 {
				t.DownloadAll()
			} else {
				for _, f := range t.Files() {
					for _, fileArg := range flags.File {
						if f.DisplayPath() == fileArg {
							f.Download()
						}
					}
				}
			}
		}()
	}
	return nil
}

var flags struct {
	Debug bool

	*DownloadCmd      `arg:"subcommand:download"`
	*ListFilesCmd     `arg:"subcommand:list-files"`
	*SpewBencodingCmd `arg:"subcommand:spew-bencoding"`
	*AnnounceCmd      `arg:"subcommand:announce"`
}

type SpewBencodingCmd struct{}

type DownloadCmd struct {
	Mmap            bool           `help:"memory-map torrent data"`
	TestPeer        []string       `help:"addresses of some starting peers"`
	Seed            bool           `help:"seed after download is complete"`
	Addr            string         `help:"network listen addr"`
	UploadRate      *tagflag.Bytes `help:"max piece bytes to send per second"`
	DownloadRate    *tagflag.Bytes `help:"max bytes per second down from peers"`
	PackedBlocklist string
	PublicIP        net.IP
	Progress        bool `default:"true"`
	PieceStates     bool
	Quiet           bool  `help:"discard client logging"`
	Stats           *bool `help:"print stats at termination"`
	Dht             bool  `default:"true"`

	TcpPeers        bool `default:"true"`
	UtpPeers        bool `default:"true"`
	Webtorrent      bool `default:"true"`
	DisableWebseeds bool

	Ipv4 bool `default:"true"`
	Ipv6 bool `default:"true"`
	Pex  bool `default:"true"`

	File    []string
	Torrent []string `arity:"+" help:"torrent file path or magnet uri" arg:"positional"`
}

type ListFilesCmd struct {
	TorrentPath string `arg:"positional"`
}

func stdoutAndStderrAreSameFile() bool {
	fi1, _ := os.Stdout.Stat()
	fi2, _ := os.Stderr.Stat()
	return os.SameFile(fi1, fi2)
}

func statsEnabled() bool {
	if flags.Stats == nil {
		return flags.Debug
	}
	return *flags.Stats
}

func exitSignalHandlers(notify *missinggo.SynchronizedEvent) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
	for {
		log.Printf("close signal received: %+v", <-c)
		notify.Set()
	}
}

func main() {
	if err := mainErr(); err != nil {
		log.Printf("error in main: %v", err)
		os.Exit(1)
	}
}

func mainErr() error {
	stdLog.SetFlags(stdLog.Flags() | stdLog.Lshortfile)
	p := arg.MustParse(&flags)
	switch {
	case flags.AnnounceCmd != nil:
		return announceErr()
	//case :
	//	return announceErr(flags.Args, parser)
	case flags.DownloadCmd != nil:
		return downloadErr()
	case flags.ListFilesCmd != nil:
		mi, err := metainfo.LoadFromFile(flags.ListFilesCmd.TorrentPath)
		if err != nil {
			return fmt.Errorf("loading from file %q: %v", flags.ListFilesCmd.TorrentPath, err)
		}
		info, err := mi.UnmarshalInfo()
		if err != nil {
			return fmt.Errorf("unmarshalling info from metainfo at %q: %v", flags.ListFilesCmd.TorrentPath, err)
		}
		for _, f := range info.UpvertedFiles() {
			fmt.Println(f.DisplayPath(&info))
		}
		return nil
	case flags.SpewBencodingCmd != nil:
		d := bencode.NewDecoder(os.Stdin)
		for i := 0; ; i++ {
			var v interface{}
			err := d.Decode(&v)
			if err == io.EOF {
				break
			}
			if err != nil {
				return fmt.Errorf("decoding message index %d: %w", i, err)
			}
			spew.Dump(v)
		}
		return nil
	default:
		p.Fail(fmt.Sprintf("unexpected subcommand: %v", p.Subcommand()))
		panic("unreachable")
	}
}

func downloadErr() error {
	defer envpprof.Stop()
	clientConfig := torrent.NewDefaultClientConfig()
	clientConfig.DisableWebseeds = flags.DisableWebseeds
	clientConfig.DisableTCP = !flags.TcpPeers
	clientConfig.DisableUTP = !flags.UtpPeers
	clientConfig.DisableIPv4 = !flags.Ipv4
	clientConfig.DisableIPv6 = !flags.Ipv6
	clientConfig.DisableAcceptRateLimiting = true
	clientConfig.NoDHT = !flags.Dht
	clientConfig.Debug = flags.Debug
	clientConfig.Seed = flags.Seed
	clientConfig.PublicIp4 = flags.PublicIP
	clientConfig.PublicIp6 = flags.PublicIP
	clientConfig.DisablePEX = !flags.Pex
	clientConfig.DisableWebtorrent = !flags.Webtorrent
	if flags.PackedBlocklist != "" {
		blocklist, err := iplist.MMapPackedFile(flags.PackedBlocklist)
		if err != nil {
			return xerrors.Errorf("loading blocklist: %v", err)
		}
		defer blocklist.Close()
		clientConfig.IPBlocklist = blocklist
	}
	if flags.Mmap {
		clientConfig.DefaultStorage = storage.NewMMap("")
	}
	if flags.Addr != "" {
		clientConfig.SetListenAddr(flags.Addr)
	}
	if flags.UploadRate != nil {
		clientConfig.UploadRateLimiter = rate.NewLimiter(rate.Limit(*flags.UploadRate), 256<<10)
	}
	if flags.DownloadRate != nil {
		clientConfig.DownloadRateLimiter = rate.NewLimiter(rate.Limit(*flags.DownloadRate), 1<<20)
	}
	if flags.Quiet {
		clientConfig.Logger = log.Discard
	}

	var stop missinggo.SynchronizedEvent
	defer func() {
		stop.Set()
	}()

	client, err := torrent.NewClient(clientConfig)
	if err != nil {
		return xerrors.Errorf("creating client: %v", err)
	}
	defer client.Close()
	go exitSignalHandlers(&stop)
	go func() {
		<-stop.C()
		client.Close()
	}()

	// Write status on the root path on the default HTTP muxer. This will be bound to localhost
	// somewhere if GOPPROF is set, thanks to the envpprof import.
	http.HandleFunc("/", func(w http.ResponseWriter, req *http.Request) {
		client.WriteStatus(w)
	})
	err = addTorrents(client)
	if err != nil {
		return fmt.Errorf("adding torrents: %w", err)
	}
	defer outputStats(client)
	if client.WaitAll() {
		log.Print("downloaded ALL the torrents")
	} else {
		return xerrors.New("y u no complete torrents?!")
	}
	if flags.Seed {
		outputStats(client)
		<-stop.C()
	}
	return nil
}

func outputStats(cl *torrent.Client) {
	if !statsEnabled() {
		return
	}
	expvar.Do(func(kv expvar.KeyValue) {
		fmt.Printf("%s: %s\n", kv.Key, kv.Value)
	})
	cl.WriteStatus(os.Stdout)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/keybase/go-libkb"
)

type CmdMykeySelect struct {
	arg   keyGenArg
	query string
}

func (v *CmdMykeySelect) ParseArgv(ctx *cli.Context) (err error) {
	nargs := len(ctx.Args())
	if err = v.arg.ParseArgv(ctx); err != nil {
	} else if nargs == 1 {
		v.query = ctx.Args()[0]
	} else if nargs != 0 {
		err = fmt.Errorf("mkey select takes 0 or 1 arguments")
	}
	return err
}

func (v *CmdMykeySelect) Run() error {
	gpg := G.GetGpgClient()
	if _, err := gpg.Configure(); err != nil {
		return err
	}
	index, err, warns := gpg.Index(true, v.query)
	if err != nil {
		return err
	}
	warns.Warn()
	if len(index.Keys) > 1 {
		if len(v.query) > 0 {
			G_UI.Output("Multiple keys matched '" + v.query + "':\n")
		} else {
			G_UI.Output("Multiple keys found:\n")
		}
	} else {
		G.Log.Info("Key selection is unambiguous: %s", index.Keys[0].GetFingerprint().ToQuads())
	}
	return nil
}

func NewCmdMykeySelect(cl *CommandLine) cli.Command {
	return cli.Command{
		Name:        "select",
		Usage:       "keybase mykey select [<key-query>]",
		Description: "Select a key as your own and push it to the server",
		Action: func(c *cli.Context) {
			cl.ChooseCommand(&CmdMykeySelect{}, "select", c)
		},
	}
}

func (v *CmdMykeySelect) GetUsage() libkb.Usage {
	return libkb.Usage{
		Config:    true,
		KbKeyring: true,
		API:       true,
		Terminal:  true,
	}
}

var example string = `
Multiple keys were found that matched 'themax':
 (1)   4096R  6052B2AD31A6631C  exp: 2024-2-3   themax@gmail.com          max@keybase.io        max@maxk.org  krohn@mit.edu  krohn@post.harvard.edu  krohn@alum.mit.edu
 (2)   4096R  081C5490F9F7B7AB  exp: never      max891@keybase.io         themax+891@gmail.com
 (3)   1024D  2D7BA4DBE9378FD9  exp: 2024-3-11  max85@keybase.io          themax+85@gmail.com
 (4)   4096R  63847B4B83930F0C  exp: 2017-10-4  themax@gmail.com          themax@gmail.com
 (5)   2048R  B4C0C0B6EF6BA2A3  exp: never      themax+3219@gmail.com
 (6)   1024D  2FC3671F312AA0FD  exp: never      themax+942@gmail.com
 (7)   2048R  D707EE8287B62514  exp: 2024-3-31  themax+309@gmail.com
 (8)   2048R  F77B98E4D4D80001  exp: 2024-3-28  themax+105.2@gmail.com
 (9)   2048R  BC07EE4BF5A71E11  exp: 2024-3-28  themax+test105@gmail.com
 (10)  2048R  A049F8D7E155170C  exp: 2024-3-24  themax+47@gmail.com
 (11)  1600R  D1DF491E66300928  exp: 2024-3-10  themax+91@gmail.com
 (12)  1024R  36B8DA42022CDC30  exp: 2024-3-10  themax+81@gmail.com
 (13)  2048R  2EE0695E30C55BEF  exp: 2024-2-17  themax+1@gmail.com
`
<commit_msg>tablification<commit_after>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/keybase/go-libkb"
	"os"
)

type CmdMykeySelect struct {
	arg   keyGenArg
	query string
}

func (v *CmdMykeySelect) ParseArgv(ctx *cli.Context) (err error) {
	nargs := len(ctx.Args())
	if err = v.arg.ParseArgv(ctx); err != nil {
	} else if nargs == 1 {
		v.query = ctx.Args()[0]
	} else if nargs != 0 {
		err = fmt.Errorf("mkey select takes 0 or 1 arguments")
	}
	return err
}

func (v *CmdMykeySelect) Run() error {
	gpg := G.GetGpgClient()
	if _, err := gpg.Configure(); err != nil {
		return err
	}
	index, err, warns := gpg.Index(true, v.query)
	if err != nil {
		return err
	}
	warns.Warn()
	if len(index.Keys) > 1 {
		if len(v.query) > 0 {
			G_UI.Output("Multiple keys matched '" + v.query + "':\n")
		} else {
			G_UI.Output("Multiple keys found:\n")
		}
		libkb.Tablify(os.Stdout, nil, index.GetRowFunc())
	} else {
		G.Log.Info("Key selection is unambiguous: %s", index.Keys[0].GetFingerprint().ToQuads())
	}
	return nil
}

func NewCmdMykeySelect(cl *CommandLine) cli.Command {
	return cli.Command{
		Name:        "select",
		Usage:       "keybase mykey select [<key-query>]",
		Description: "Select a key as your own and push it to the server",
		Action: func(c *cli.Context) {
			cl.ChooseCommand(&CmdMykeySelect{}, "select", c)
		},
	}
}

func (v *CmdMykeySelect) GetUsage() libkb.Usage {
	return libkb.Usage{
		Config:    true,
		KbKeyring: true,
		API:       true,
		Terminal:  true,
	}
}

var example string = `
Multiple keys were found that matched 'themax':
 (1)   4096R  6052B2AD31A6631C  exp: 2024-2-3   themax@gmail.com          max@keybase.io        max@maxk.org  krohn@mit.edu  krohn@post.harvard.edu  krohn@alum.mit.edu
 (2)   4096R  081C5490F9F7B7AB  exp: never      max891@keybase.io         themax+891@gmail.com
 (3)   1024D  2D7BA4DBE9378FD9  exp: 2024-3-11  max85@keybase.io          themax+85@gmail.com
 (4)   4096R  63847B4B83930F0C  exp: 2017-10-4  themax@gmail.com          themax@gmail.com
 (5)   2048R  B4C0C0B6EF6BA2A3  exp: never      themax+3219@gmail.com
 (6)   1024D  2FC3671F312AA0FD  exp: never      themax+942@gmail.com
 (7)   2048R  D707EE8287B62514  exp: 2024-3-31  themax+309@gmail.com
 (8)   2048R  F77B98E4D4D80001  exp: 2024-3-28  themax+105.2@gmail.com
 (9)   2048R  BC07EE4BF5A71E11  exp: 2024-3-28  themax+test105@gmail.com
 (10)  2048R  A049F8D7E155170C  exp: 2024-3-24  themax+47@gmail.com
 (11)  1600R  D1DF491E66300928  exp: 2024-3-10  themax+91@gmail.com
 (12)  1024R  36B8DA42022CDC30  exp: 2024-3-10  themax+81@gmail.com
 (13)  2048R  2EE0695E30C55BEF  exp: 2024-2-17  themax+1@gmail.com
`
<|endoftext|>


<commit_before>// Copyright 2019-present PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package cophandler

import (
	"math"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/unistore/tikv/dbreader"
	"github.com/pingcap/badger/y"
	"github.com/pingcap/kvproto/pkg/coprocessor"
	"github.com/pingcap/parser/ast"
	"github.com/pingcap/parser/charset"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/parser/mysql"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/statistics"
	"github.com/pingcap/tidb/tablecodec"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"github.com/pingcap/tidb/util/collate"
	"github.com/pingcap/tidb/util/rowcodec"
	"github.com/pingcap/tipb/go-tipb"
	"golang.org/x/net/context"
)

// handleCopAnalyzeRequest handles coprocessor analyze request.
func handleCopAnalyzeRequest(dbReader *dbreader.DBReader, req *coprocessor.Request) *coprocessor.Response {
	resp := &coprocessor.Response{}
	if len(req.Ranges) == 0 {
		return resp
	}
	if req.GetTp() != kv.ReqTypeAnalyze {
		return resp
	}
	analyzeReq := new(tipb.AnalyzeReq)
	err := proto.Unmarshal(req.Data, analyzeReq)
	if err != nil {
		resp.OtherError = err.Error()
		return resp
	}
	ranges, err := extractKVRanges(dbReader.StartKey, dbReader.EndKey, req.Ranges, false)
	if err != nil {
		resp.OtherError = err.Error()
		return resp
	}
	y.Assert(len(ranges) == 1)
	if analyzeReq.Tp == tipb.AnalyzeType_TypeIndex {
		resp, err = handleAnalyzeIndexReq(dbReader, ranges[0], analyzeReq, req.StartTs)
	} else if analyzeReq.Tp == tipb.AnalyzeType_TypeCommonHandle {
		resp, err = handleAnalyzeCommonHandleReq(dbReader, ranges[0], analyzeReq, req.StartTs)
	} else {
		resp, err = handleAnalyzeColumnsReq(dbReader, ranges[0], analyzeReq, req.StartTs)
	}
	if err != nil {
		resp = &coprocessor.Response{
			OtherError: err.Error(),
		}
	}
	return resp
}

func handleAnalyzeIndexReq(dbReader *dbreader.DBReader, ran kv.KeyRange, analyzeReq *tipb.AnalyzeReq, startTS uint64) (*coprocessor.Response, error) {
	processor := &analyzeIndexProcessor{
		colLen:       int(analyzeReq.IdxReq.NumColumns),
		statsBuilder: statistics.NewSortedBuilder(flagsToStatementContext(analyzeReq.Flags), analyzeReq.IdxReq.BucketSize, 0, types.NewFieldType(mysql.TypeBlob)),
	}
	if analyzeReq.IdxReq.CmsketchDepth != nil && analyzeReq.IdxReq.CmsketchWidth != nil {
		processor.cms = statistics.NewCMSketch(*analyzeReq.IdxReq.CmsketchDepth, *analyzeReq.IdxReq.CmsketchWidth)
	}
	err := dbReader.Scan(ran.StartKey, ran.EndKey, math.MaxInt64, startTS, processor)
	if err != nil {
		return nil, err
	}
	hg := statistics.HistogramToProto(processor.statsBuilder.Hist())
	var cm *tipb.CMSketch
	if processor.cms != nil {
		cm = statistics.CMSketchToProto(processor.cms)
	}
	data, err := proto.Marshal(&tipb.AnalyzeIndexResp{Hist: hg, Cms: cm})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return &coprocessor.Response{Data: data}, nil
}

func handleAnalyzeCommonHandleReq(dbReader *dbreader.DBReader, ran kv.KeyRange, analyzeReq *tipb.AnalyzeReq, startTS uint64) (*coprocessor.Response, error) {
	processor := &analyzeCommonHandleProcessor{
		colLen:       int(analyzeReq.IdxReq.NumColumns),
		statsBuilder: statistics.NewSortedBuilder(flagsToStatementContext(analyzeReq.Flags), analyzeReq.IdxReq.BucketSize, 0, types.NewFieldType(mysql.TypeBlob)),
	}
	if analyzeReq.IdxReq.CmsketchDepth != nil && analyzeReq.IdxReq.CmsketchWidth != nil {
		processor.cms = statistics.NewCMSketch(*analyzeReq.IdxReq.CmsketchDepth, *analyzeReq.IdxReq.CmsketchWidth)
	}
	err := dbReader.Scan(ran.StartKey, ran.EndKey, math.MaxInt64, startTS, processor)
	if err != nil {
		return nil, err
	}
	hg := statistics.HistogramToProto(processor.statsBuilder.Hist())
	var cm *tipb.CMSketch
	if processor.cms != nil {
		cm = statistics.CMSketchToProto(processor.cms)
	}
	data, err := proto.Marshal(&tipb.AnalyzeIndexResp{Hist: hg, Cms: cm})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return &coprocessor.Response{Data: data}, nil
}

type analyzeIndexProcessor struct {
	skipVal

	colLen       int
	statsBuilder *statistics.SortedBuilder
	cms          *statistics.CMSketch
	rowBuf       []byte
}

func (p *analyzeIndexProcessor) Process(key, value []byte) error {
	values, _, err := tablecodec.CutIndexKeyNew(key, p.colLen)
	if err != nil {
		return err
	}
	p.rowBuf = p.rowBuf[:0]
	for _, val := range values {
		p.rowBuf = append(p.rowBuf, val...)
		if p.cms != nil {
			p.cms.InsertBytes(p.rowBuf)
		}
	}
	rowData := safeCopy(p.rowBuf)
	err = p.statsBuilder.Iterate(types.NewBytesDatum(rowData))
	if err != nil {
		return err
	}
	return nil
}

type analyzeCommonHandleProcessor struct {
	skipVal

	colLen       int
	statsBuilder *statistics.SortedBuilder
	cms          *statistics.CMSketch
	rowBuf       []byte
}

func (p *analyzeCommonHandleProcessor) Process(key, value []byte) error {
	values, _, err := tablecodec.CutCommonHandle(key, p.colLen)
	if err != nil {
		return err
	}
	p.rowBuf = p.rowBuf[:0]
	for _, val := range values {
		p.rowBuf = append(p.rowBuf, val...)
		if p.cms != nil {
			p.cms.InsertBytes(p.rowBuf)
		}
	}
	rowData := safeCopy(p.rowBuf)
	err = p.statsBuilder.Iterate(types.NewBytesDatum(rowData))
	if err != nil {
		return err
	}
	return nil
}

type analyzeColumnsExec struct {
	skipVal
	reader  *dbreader.DBReader
	seekKey []byte
	endKey  []byte
	startTS uint64

	chk     *chunk.Chunk
	decoder *rowcodec.ChunkDecoder
	req     *chunk.Chunk
	evalCtx *evalContext
	fields  []*ast.ResultField
}

func handleAnalyzeColumnsReq(dbReader *dbreader.DBReader, ran kv.KeyRange, analyzeReq *tipb.AnalyzeReq, startTS uint64) (*coprocessor.Response, error) {
	sc := flagsToStatementContext(analyzeReq.Flags)
	sc.TimeZone = time.FixedZone("UTC", int(analyzeReq.TimeZoneOffset))
	evalCtx := &evalContext{sc: sc}
	columns := analyzeReq.ColReq.ColumnsInfo
	evalCtx.setColumnInfo(columns)
	if len(analyzeReq.ColReq.PrimaryColumnIds) > 0 {
		evalCtx.primaryCols = analyzeReq.ColReq.PrimaryColumnIds
	}
	decoder, err := evalCtx.newRowDecoder()
	if err != nil {
		return nil, err
	}
	e := &analyzeColumnsExec{
		reader:  dbReader,
		seekKey: ran.StartKey,
		endKey:  ran.EndKey,
		startTS: startTS,
		chk:     chunk.NewChunkWithCapacity(evalCtx.fieldTps, 1),
		decoder: decoder,
		evalCtx: evalCtx,
	}
	e.fields = make([]*ast.ResultField, len(columns))
	for i := range e.fields {
		rf := new(ast.ResultField)
		rf.Column = new(model.ColumnInfo)
		rf.Column.FieldType = types.FieldType{Tp: mysql.TypeBlob, Flen: mysql.MaxBlobWidth, Charset: charset.CharsetUTF8, Collate: charset.CollationUTF8}
		e.fields[i] = rf
	}

	pkID := int64(-1)
	numCols := len(columns)
	if columns[0].GetPkHandle() {
		pkID = columns[0].ColumnId
		columns = columns[1:]
		numCols--
	}
	collators := make([]collate.Collator, numCols)
	fts := make([]*types.FieldType, numCols)
	for i, col := range columns {
		ft := fieldTypeFromPBColumn(col)
		fts[i] = ft
		if ft.EvalType() == types.ETString {
			collators[i] = collate.GetCollator(ft.Collate)
		}
	}
	colReq := analyzeReq.ColReq
	builder := statistics.SampleBuilder{
		Sc:              sc,
		RecordSet:       e,
		ColLen:          numCols,
		MaxBucketSize:   colReq.BucketSize,
		MaxFMSketchSize: colReq.SketchSize,
		MaxSampleSize:   colReq.SampleSize,
		Collators:       collators,
		ColsFieldType:   fts,
	}
	if pkID != -1 {
		builder.PkBuilder = statistics.NewSortedBuilder(sc, builder.MaxBucketSize, pkID, types.NewFieldType(mysql.TypeBlob))
	}
	if colReq.CmsketchWidth != nil && colReq.CmsketchDepth != nil {
		builder.CMSketchWidth = *colReq.CmsketchWidth
		builder.CMSketchDepth = *colReq.CmsketchDepth
	}
	collectors, pkBuilder, err := builder.CollectColumnStats()
	if err != nil {
		return nil, errors.Trace(err)
	}
	colResp := &tipb.AnalyzeColumnsResp{}
	if pkID != -1 {
		colResp.PkHist = statistics.HistogramToProto(pkBuilder.Hist())
	}
	for _, c := range collectors {
		colResp.Collectors = append(colResp.Collectors, statistics.SampleCollectorToProto(c))
	}
	data, err := proto.Marshal(colResp)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return &coprocessor.Response{Data: data}, nil
}

// Fields implements the sqlexec.RecordSet Fields interface.
func (e *analyzeColumnsExec) Fields() []*ast.ResultField {
	return e.fields
}

func (e *analyzeColumnsExec) Next(ctx context.Context, req *chunk.Chunk) error {
	req.Reset()
	e.req = req
	err := e.reader.Scan(e.seekKey, e.endKey, math.MaxInt64, e.startTS, e)
	if err != nil {
		return err
	}
	if req.NumRows() < req.Capacity() {
		e.seekKey = e.endKey
	}
	return nil
}

func (e *analyzeColumnsExec) Process(key, value []byte) error {
	handle, err := tablecodec.DecodeRowKey(key)
	if err != nil {
		return errors.Trace(err)
	}
	err = e.decoder.DecodeToChunk(value, handle, e.chk)
	if err != nil {
		return errors.Trace(err)
	}
	row := e.chk.GetRow(0)
	for i, tp := range e.evalCtx.fieldTps {
		d := row.GetDatum(i, tp)
		if d.IsNull() {
			e.req.AppendNull(i)
			continue
		}

		value, err := tablecodec.EncodeValue(e.evalCtx.sc, nil, d)
		if err != nil {
			return err
		}
		e.req.AppendBytes(i, value)
	}
	e.chk.Reset()
	if e.req.NumRows() == e.req.Capacity() {
		e.seekKey = kv.Key(key).PrefixNext()
		return dbreader.ScanBreak
	}
	return nil
}

func (e *analyzeColumnsExec) NewChunk() *chunk.Chunk {
	fields := make([]*types.FieldType, 0, len(e.fields))
	for _, field := range e.fields {
		fields = append(fields, &field.Column.FieldType)
	}
	return chunk.NewChunkWithCapacity(fields, 1024)
}

// Close implements the sqlexec.RecordSet Close interface.
func (e *analyzeColumnsExec) Close() error {
	return nil
}
<commit_msg>unistore: support analyze multi-ranges (#19838)<commit_after>// Copyright 2019-present PingCAP, 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package cophandler

import (
	"math"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/juju/errors"
	"github.com/ngaut/unistore/tikv/dbreader"
	"github.com/pingcap/badger/y"
	"github.com/pingcap/kvproto/pkg/coprocessor"
	"github.com/pingcap/parser/ast"
	"github.com/pingcap/parser/charset"
	"github.com/pingcap/parser/model"
	"github.com/pingcap/parser/mysql"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/statistics"
	"github.com/pingcap/tidb/tablecodec"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"github.com/pingcap/tidb/util/collate"
	"github.com/pingcap/tidb/util/rowcodec"
	"github.com/pingcap/tipb/go-tipb"
	"golang.org/x/net/context"
)

// handleCopAnalyzeRequest handles coprocessor analyze request.
func handleCopAnalyzeRequest(dbReader *dbreader.DBReader, req *coprocessor.Request) *coprocessor.Response {
	resp := &coprocessor.Response{}
	if len(req.Ranges) == 0 {
		return resp
	}
	if req.GetTp() != kv.ReqTypeAnalyze {
		return resp
	}
	analyzeReq := new(tipb.AnalyzeReq)
	err := proto.Unmarshal(req.Data, analyzeReq)
	if err != nil {
		resp.OtherError = err.Error()
		return resp
	}
	ranges, err := extractKVRanges(dbReader.StartKey, dbReader.EndKey, req.Ranges, false)
	if err != nil {
		resp.OtherError = err.Error()
		return resp
	}
	y.Assert(len(ranges) >= 1)
	if analyzeReq.Tp == tipb.AnalyzeType_TypeIndex {
		resp, err = handleAnalyzeIndexReq(dbReader, ranges, analyzeReq, req.StartTs)
	} else if analyzeReq.Tp == tipb.AnalyzeType_TypeCommonHandle {
		resp, err = handleAnalyzeCommonHandleReq(dbReader, ranges, analyzeReq, req.StartTs)
	} else {
		resp, err = handleAnalyzeColumnsReq(dbReader, ranges, analyzeReq, req.StartTs)
	}
	if err != nil {
		resp = &coprocessor.Response{
			OtherError: err.Error(),
		}
	}
	return resp
}

func handleAnalyzeIndexReq(dbReader *dbreader.DBReader, rans []kv.KeyRange, analyzeReq *tipb.AnalyzeReq, startTS uint64) (*coprocessor.Response, error) {
	processor := &analyzeIndexProcessor{
		colLen:       int(analyzeReq.IdxReq.NumColumns),
		statsBuilder: statistics.NewSortedBuilder(flagsToStatementContext(analyzeReq.Flags), analyzeReq.IdxReq.BucketSize, 0, types.NewFieldType(mysql.TypeBlob)),
	}
	if analyzeReq.IdxReq.CmsketchDepth != nil && analyzeReq.IdxReq.CmsketchWidth != nil {
		processor.cms = statistics.NewCMSketch(*analyzeReq.IdxReq.CmsketchDepth, *analyzeReq.IdxReq.CmsketchWidth)
	}
	for _, ran := range rans {
		err := dbReader.Scan(ran.StartKey, ran.EndKey, math.MaxInt64, startTS, processor)
		if err != nil {
			return nil, err
		}
	}
	hg := statistics.HistogramToProto(processor.statsBuilder.Hist())
	var cm *tipb.CMSketch
	if processor.cms != nil {
		cm = statistics.CMSketchToProto(processor.cms)
	}
	data, err := proto.Marshal(&tipb.AnalyzeIndexResp{Hist: hg, Cms: cm})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return &coprocessor.Response{Data: data}, nil
}

func handleAnalyzeCommonHandleReq(dbReader *dbreader.DBReader, rans []kv.KeyRange, analyzeReq *tipb.AnalyzeReq, startTS uint64) (*coprocessor.Response, error) {
	processor := &analyzeCommonHandleProcessor{
		colLen:       int(analyzeReq.IdxReq.NumColumns),
		statsBuilder: statistics.NewSortedBuilder(flagsToStatementContext(analyzeReq.Flags), analyzeReq.IdxReq.BucketSize, 0, types.NewFieldType(mysql.TypeBlob)),
	}
	if analyzeReq.IdxReq.CmsketchDepth != nil && analyzeReq.IdxReq.CmsketchWidth != nil {
		processor.cms = statistics.NewCMSketch(*analyzeReq.IdxReq.CmsketchDepth, *analyzeReq.IdxReq.CmsketchWidth)
	}
	for _, ran := range rans {
		err := dbReader.Scan(ran.StartKey, ran.EndKey, math.MaxInt64, startTS, processor)
		if err != nil {
			return nil, err
		}
	}
	hg := statistics.HistogramToProto(processor.statsBuilder.Hist())
	var cm *tipb.CMSketch
	if processor.cms != nil {
		cm = statistics.CMSketchToProto(processor.cms)
	}
	data, err := proto.Marshal(&tipb.AnalyzeIndexResp{Hist: hg, Cms: cm})
	if err != nil {
		return nil, errors.Trace(err)
	}
	return &coprocessor.Response{Data: data}, nil
}

type analyzeIndexProcessor struct {
	skipVal

	colLen       int
	statsBuilder *statistics.SortedBuilder
	cms          *statistics.CMSketch
	rowBuf       []byte
}

func (p *analyzeIndexProcessor) Process(key, value []byte) error {
	values, _, err := tablecodec.CutIndexKeyNew(key, p.colLen)
	if err != nil {
		return err
	}
	p.rowBuf = p.rowBuf[:0]
	for _, val := range values {
		p.rowBuf = append(p.rowBuf, val...)
		if p.cms != nil {
			p.cms.InsertBytes(p.rowBuf)
		}
	}
	rowData := safeCopy(p.rowBuf)
	err = p.statsBuilder.Iterate(types.NewBytesDatum(rowData))
	if err != nil {
		return err
	}
	return nil
}

type analyzeCommonHandleProcessor struct {
	skipVal

	colLen       int
	statsBuilder *statistics.SortedBuilder
	cms          *statistics.CMSketch
	rowBuf       []byte
}

func (p *analyzeCommonHandleProcessor) Process(key, value []byte) error {
	values, _, err := tablecodec.CutCommonHandle(key, p.colLen)
	if err != nil {
		return err
	}
	p.rowBuf = p.rowBuf[:0]
	for _, val := range values {
		p.rowBuf = append(p.rowBuf, val...)
		if p.cms != nil {
			p.cms.InsertBytes(p.rowBuf)
		}
	}
	rowData := safeCopy(p.rowBuf)
	err = p.statsBuilder.Iterate(types.NewBytesDatum(rowData))
	if err != nil {
		return err
	}
	return nil
}

type analyzeColumnsExec struct {
	skipVal
	reader  *dbreader.DBReader
	ranges  []kv.KeyRange
	curRan  int
	seekKey []byte
	endKey  []byte
	startTS uint64

	chk     *chunk.Chunk
	decoder *rowcodec.ChunkDecoder
	req     *chunk.Chunk
	evalCtx *evalContext
	fields  []*ast.ResultField
}

func handleAnalyzeColumnsReq(dbReader *dbreader.DBReader, rans []kv.KeyRange, analyzeReq *tipb.AnalyzeReq, startTS uint64) (*coprocessor.Response, error) {
	sc := flagsToStatementContext(analyzeReq.Flags)
	sc.TimeZone = time.FixedZone("UTC", int(analyzeReq.TimeZoneOffset))
	evalCtx := &evalContext{sc: sc}
	columns := analyzeReq.ColReq.ColumnsInfo
	evalCtx.setColumnInfo(columns)
	if len(analyzeReq.ColReq.PrimaryColumnIds) > 0 {
		evalCtx.primaryCols = analyzeReq.ColReq.PrimaryColumnIds
	}
	decoder, err := evalCtx.newRowDecoder()
	if err != nil {
		return nil, err
	}
	e := &analyzeColumnsExec{
		reader:  dbReader,
		seekKey: rans[0].StartKey,
		endKey:  rans[0].EndKey,
		ranges:  rans,
		curRan:  0,
		startTS: startTS,
		chk:     chunk.NewChunkWithCapacity(evalCtx.fieldTps, 1),
		decoder: decoder,
		evalCtx: evalCtx,
	}
	e.fields = make([]*ast.ResultField, len(columns))
	for i := range e.fields {
		rf := new(ast.ResultField)
		rf.Column = new(model.ColumnInfo)
		rf.Column.FieldType = types.FieldType{Tp: mysql.TypeBlob, Flen: mysql.MaxBlobWidth, Charset: charset.CharsetUTF8, Collate: charset.CollationUTF8}
		e.fields[i] = rf
	}

	pkID := int64(-1)
	numCols := len(columns)
	if columns[0].GetPkHandle() {
		pkID = columns[0].ColumnId
		columns = columns[1:]
		numCols--
	}
	collators := make([]collate.Collator, numCols)
	fts := make([]*types.FieldType, numCols)
	for i, col := range columns {
		ft := fieldTypeFromPBColumn(col)
		fts[i] = ft
		if ft.EvalType() == types.ETString {
			collators[i] = collate.GetCollator(ft.Collate)
		}
	}
	colReq := analyzeReq.ColReq
	builder := statistics.SampleBuilder{
		Sc:              sc,
		RecordSet:       e,
		ColLen:          numCols,
		MaxBucketSize:   colReq.BucketSize,
		MaxFMSketchSize: colReq.SketchSize,
		MaxSampleSize:   colReq.SampleSize,
		Collators:       collators,
		ColsFieldType:   fts,
	}
	if pkID != -1 {
		builder.PkBuilder = statistics.NewSortedBuilder(sc, builder.MaxBucketSize, pkID, types.NewFieldType(mysql.TypeBlob))
	}
	if colReq.CmsketchWidth != nil && colReq.CmsketchDepth != nil {
		builder.CMSketchWidth = *colReq.CmsketchWidth
		builder.CMSketchDepth = *colReq.CmsketchDepth
	}
	collectors, pkBuilder, err := builder.CollectColumnStats()
	if err != nil {
		return nil, errors.Trace(err)
	}
	colResp := &tipb.AnalyzeColumnsResp{}
	if pkID != -1 {
		colResp.PkHist = statistics.HistogramToProto(pkBuilder.Hist())
	}
	for _, c := range collectors {
		colResp.Collectors = append(colResp.Collectors, statistics.SampleCollectorToProto(c))
	}
	data, err := proto.Marshal(colResp)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return &coprocessor.Response{Data: data}, nil
}

// Fields implements the sqlexec.RecordSet Fields interface.
func (e *analyzeColumnsExec) Fields() []*ast.ResultField {
	return e.fields
}

func (e *analyzeColumnsExec) Next(ctx context.Context, req *chunk.Chunk) error {
	req.Reset()
	e.req = req
	err := e.reader.Scan(e.seekKey, e.endKey, math.MaxInt64, e.startTS, e)
	if err != nil {
		return err
	}
	if req.NumRows() < req.Capacity() {
		if e.curRan == len(e.ranges)-1 {
			e.seekKey = e.endKey
		} else {
			e.curRan++
			e.seekKey = e.ranges[e.curRan].StartKey
			e.endKey = e.ranges[e.curRan].EndKey
		}
	}
	return nil
}

func (e *analyzeColumnsExec) Process(key, value []byte) error {
	handle, err := tablecodec.DecodeRowKey(key)
	if err != nil {
		return errors.Trace(err)
	}
	err = e.decoder.DecodeToChunk(value, handle, e.chk)
	if err != nil {
		return errors.Trace(err)
	}
	row := e.chk.GetRow(0)
	for i, tp := range e.evalCtx.fieldTps {
		d := row.GetDatum(i, tp)
		if d.IsNull() {
			e.req.AppendNull(i)
			continue
		}

		value, err := tablecodec.EncodeValue(e.evalCtx.sc, nil, d)
		if err != nil {
			return err
		}
		e.req.AppendBytes(i, value)
	}
	e.chk.Reset()
	if e.req.NumRows() == e.req.Capacity() {
		e.seekKey = kv.Key(key).PrefixNext()
		return dbreader.ScanBreak
	}
	return nil
}

func (e *analyzeColumnsExec) NewChunk() *chunk.Chunk {
	fields := make([]*types.FieldType, 0, len(e.fields))
	for _, field := range e.fields {
		fields = append(fields, &field.Column.FieldType)
	}
	return chunk.NewChunkWithCapacity(fields, 1024)
}

// Close implements the sqlexec.RecordSet Close interface.
func (e *analyzeColumnsExec) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package routes

import (
	"net/http"
	"regexp"
	"strings"

	"github.com/waltzofpearls/tetris-go/libs"
)

type Index struct {
	libs.Subrouter
}

func (sr *Index) AttachHandlers() {
	sr.Router.Handle("/", libs.HandlerFunc(sr.homeHandler)).Methods("GET")
	sr.Router.Handle("/{redirect:(projects|resume)}", libs.HandlerFunc(sr.redirectHandler)).Methods("GET")
}

func (sr *Index) homeHandler(w http.ResponseWriter, r *http.Request) error {
	return sr.HtmlResponseHandler(w, r, "index", map[string]interface{}{
		"Config": sr.Config,
		"Title":  "Rollie Ma - Polyglot Developer from Vancouver, BC",
		"Description": strings.Join([]string{
			"Hi, I'm Rollie Ma. A Linux lover and LEGO bricks enthusiast.",
			"A polyglot developer obsessed with PHP, Golang, Python and JavaScript.",
			"A receptive learner captivated by mobile development, NoSQL and data mining.",
			"An amateur explorer interested in information aggregation and artificial intelligence fields.",
		}, " "),
		"Url":   "http://rolli3.net",
		"Image": "http://rolli3.net/images/logos/logo-120x120.png",
	})
}

func (sr *Index) redirectHandler(w http.ResponseWriter, r *http.Request) error {
	re := regexp.MustCompile("^/?")
	url := re.ReplaceAllLiteralString(r.URL.String(), "/#")
	return sr.RedirectHandler(w, r, url)
}
<commit_msg>Swap the order for my most favorite languages<commit_after>package routes

import (
	"net/http"
	"regexp"
	"strings"

	"github.com/waltzofpearls/tetris-go/libs"
)

type Index struct {
	libs.Subrouter
}

func (sr *Index) AttachHandlers() {
	sr.Router.Handle("/", libs.HandlerFunc(sr.homeHandler)).Methods("GET")
	sr.Router.Handle("/{redirect:(projects|resume)}", libs.HandlerFunc(sr.redirectHandler)).Methods("GET")
}

func (sr *Index) homeHandler(w http.ResponseWriter, r *http.Request) error {
	return sr.HtmlResponseHandler(w, r, "index", map[string]interface{}{
		"Config": sr.Config,
		"Title":  "Rollie Ma - Polyglot Developer from Vancouver, BC",
		"Description": strings.Join([]string{
			"Hi, I'm Rollie Ma. A Linux lover and LEGO bricks enthusiast.",
			"A polyglot developer obsessed with Golang, JavaScript, Python and PHP.",
			"A receptive learner captivated by mobile development, NoSQL and data mining.",
			"An amateur explorer interested in information aggregation and artificial intelligence fields.",
		}, " "),
		"Url":   "http://rolli3.net",
		"Image": "http://rolli3.net/images/logos/logo-120x120.png",
	})
}

func (sr *Index) redirectHandler(w http.ResponseWriter, r *http.Request) error {
	re := regexp.MustCompile("^/?")
	url := re.ReplaceAllLiteralString(r.URL.String(), "/#")
	return sr.RedirectHandler(w, r, url)
}
<|endoftext|>


<commit_before>package ipv4

import (
	"network/ipv4/ipv4tps"
	"time"

	"github.com/hsheth2/logs"
	"bytes"
)

type IP_Reader struct {
	incomingPackets chan []byte
	irm             *IP_Read_Manager
	protocol        uint8
	ip              *ipv4tps.IPaddress
	fragBuf         map[string](chan []byte)
}

func NewIP_Reader(irm *IP_Read_Manager, ip *ipv4tps.IPaddress, protocol uint8) (*IP_Reader, error) {
	c, err := irm.Bind(ip, protocol)
	if err != nil {
		return nil, err
	}

	return &IP_Reader{
		incomingPackets: c,
		protocol:        protocol,
		ip:              ip,
		fragBuf:         make(map[string](chan []byte)),
	}, nil
}

func slicePacket(b []byte) (hrd, payload []byte) {
	hdrLen := int(b[0]&0x0f) * 4
	//fmt.Println("HdrLen: ", hdrLen)
	return b[:hdrLen], b[hdrLen:]
}

func (ipr *IP_Reader) fragmentAssembler(in <-chan []byte, quit <-chan bool, didQuit chan<- bool, done chan bool) {
	payload := make([]byte, 0)
	extraFrags := make(map[uint64]([]byte))
	recvLast := false

	for {
		select {
		case <-quit:
			//Trace.Println("quitting upon quit signal")
			didQuit <- true
			return
		case frag := <-in:
			//Trace.Println("got a fragment packet. len:", len(frag))
			hdr, p := slicePacket(frag)
			//offset := 8 * (uint64(hdr[6]&0x1f)<<8 + uint64(hdr[7]))
			//fmt.Println("RECEIVED FRAG")
			//fmt.Println("Offset:", offset)
			//fmt.Println(len(payload))

			// add to map
			offset := 8 * (uint64(hdr[6]&0x1F)<<8 + uint64(hdr[7]))
			//Trace.Println("Offset:", offset)
			extraFrags[offset] = p

			// check more fragments flag
			if (hdr[6]>>5)&0x01 == 0 {
				recvLast = true
			}

			// add to payload
			//Trace.Println("Begin to add to the payload")
			for {
				if storedFrag, found := extraFrags[uint64(len(payload))]; found {
					//Trace.Println("New Payload Len: ", len(payload))
					delete(extraFrags, uint64(len(payload)))
					payload = append(payload, storedFrag...)
				} else {
					break
				}
			}
			//Trace.Println("Finished add to the payload")

			// deal with the payload
			if recvLast && len(extraFrags) == 0 {
				//Trace.Println("Done")
				// correct the header
				fullPacketHdr := hdr
				totalLen := uint16(fullPacketHdr[0]&0x0F)*4 + uint16(len(payload))
				fullPacketHdr[2] = byte(totalLen >> 8)
				fullPacketHdr[3] = byte(totalLen)
				fullPacketHdr[6] = 0
				fullPacketHdr[7] = 0

				// update the checksum
				check := calculateIPChecksum(fullPacketHdr[:20])
				fullPacketHdr[10] = byte(check >> 8)
				fullPacketHdr[11] = byte(check)

				// send the packet back into processing
				//go func() {
				select {
				case ipr.incomingPackets <- append(fullPacketHdr, payload...):
				default:
					logs.Warn.Println("Dropping defragmented packet, no space in buffer")
				}
				//fmt.Println("FINISHED")
				//}()
				//Trace.Println("Just wrote back in")
				done <- true
				return // from goroutine
			}
			//Trace.Println("Looping")
		}
	}

	// drop the packet upon timeout
	//Trace.Println(errors.New("Fragments took too long, packet dropped"))
	//return
}

func (ipr *IP_Reader) killFragmentAssembler(quit chan<- bool, didQuit <-chan bool, done <-chan bool, bufID string) {
	// sends quit to the assembler if it doesn't send done
	select {
	case <-time.After(FRAGMENT_TIMEOUT):
		//Trace.Println("Force quitting packet assembler")
		quit <- true
		<-didQuit // will block until it has been received
	case <-done:
		//Trace.Println("Recieved done msg.")
	}

	//Trace.Println("Frag Assemble Ended, finished")
	delete(ipr.fragBuf, bufID)
}

func (ipr *IP_Reader) ReadFrom() (rip, lip *ipv4tps.IPaddress, b, payload []byte, e error) {
	//fmt.Println("STARTING READ")
	b = <-ipr.incomingPackets
	//	logs.Info.Println("Read IP packet")
	//fmt.Println("RAW READ COMPLETED")
	//fmt.Println("Read Length: ", len(b))
	//fmt.Print(".")
	//fmt.Println("Full Read Data: ", b)

	hdr, p := slicePacket(b)

	// extract source IP and protocol
	rip = &ipv4tps.IPaddress{IP: hdr[12:16]}
	lip = &ipv4tps.IPaddress{IP: hdr[16:20]}
	proto := uint8(hdr[9])
	if !((bytes.Equal(ipr.ip.IP, rip.IP) || bytes.Equal(ipr.ip.IP, ipv4tps.IP_ALL)) && ipr.protocol == proto) {
		//Info.Println("Not interested in packet: dropping.")
		return ipr.ReadFrom()
	}

	// verify checksum
	if !verifyIPChecksum(hdr) {
		//Info.Println("Header checksum incorrect, packet dropped")
		return ipr.ReadFrom() // return another packet instead
	}

	packetOffset := uint16(hdr[6]&0x1f)<<8 + uint16(hdr[7])
	//fmt.Println("PACK OFF", packetOffset, "HEADER FLAGS", (hdr[6] >> 5))
	if ((hdr[6]>>5)&0x01 == 0) && (packetOffset == 0) {
		// not a fragment
		return rip, lip, b, p, nil
	} else {
		// is a fragment
		bufID := string([]byte{hdr[12], hdr[13], hdr[14], hdr[15], // the source IP
			hdr[16], hdr[17], hdr[18], hdr[19], // the destination IP
			hdr[9],         // protocol
			hdr[4], hdr[5], // identification
		})
		//Trace.Printf("rcvd a fragment-bufId: %x, len: %d\n", bufID, len(b))

		if _, ok := ipr.fragBuf[bufID]; !ok {
			// create the fragment buffer and quit
			//Trace.Printf("creating a new buffer for %x\n", bufID)
			ipr.fragBuf[bufID] = make(chan []byte, FRAGMENT_ASSEMBLER_BUFFER_SIZE)
			quit := make(chan bool, 1)
			done := make(chan bool, 1)
			didQuit := make(chan bool, 1)

			// create the packet assembler in a goroutine to allow the program to continue
			go ipr.fragmentAssembler(ipr.fragBuf[bufID], quit, didQuit, done)
			go ipr.killFragmentAssembler(quit, didQuit, done, bufID)
		}

		// send the packet to the assembler
		go func() { ipr.fragBuf[bufID] <- b }()

		// after dealing with the fragment, try reading again
		return ipr.ReadFrom()
	}
}

func (ipr *IP_Reader) Close() error {
	return ipr.irm.Unbind(ipr.ip, ipr.protocol)
}

/* h := &ipv4.Header{
	Version:  ipv4.Version,      // protocol version
	Len:      20,                // header length
	TOS:      0,                 // type-of-service (0 is everything normal)
	TotalLen: len(x) + 20,       // packet total length (octets)
	ID:       0,                 // identification
	Flags:    ipv4.DontFragment, // flags
	FragOff:  0,                 // fragment offset
	TTL:      8,                 // time-to-live (maximum lifespan in seconds)
	Protocol: 17,                // next protocol (17 is UDP)
	Checksum: 0,                 // checksum (apparently autocomputed)
	//Src:    net.IPv4(127, 0, 0, 1), // source address, apparently done automatically
	Dst: net.ParseIP(c.manager.ipAddress), // destination address
	//Options                         // options, extension headers
}
*/
<commit_msg>Formatted code<commit_after>package ipv4

import (
	"network/ipv4/ipv4tps"
	"time"

	"bytes"

	"github.com/hsheth2/logs"
)

type IP_Reader struct {
	incomingPackets chan []byte
	irm             *IP_Read_Manager
	protocol        uint8
	ip              *ipv4tps.IPaddress
	fragBuf         map[string](chan []byte)
}

func NewIP_Reader(irm *IP_Read_Manager, ip *ipv4tps.IPaddress, protocol uint8) (*IP_Reader, error) {
	c, err := irm.Bind(ip, protocol)
	if err != nil {
		return nil, err
	}

	return &IP_Reader{
		incomingPackets: c,
		protocol:        protocol,
		ip:              ip,
		fragBuf:         make(map[string](chan []byte)),
	}, nil
}

func slicePacket(b []byte) (hrd, payload []byte) {
	hdrLen := int(b[0]&0x0f) * 4
	//fmt.Println("HdrLen: ", hdrLen)
	return b[:hdrLen], b[hdrLen:]
}

func (ipr *IP_Reader) fragmentAssembler(in <-chan []byte, quit <-chan bool, didQuit chan<- bool, done chan bool) {
	payload := make([]byte, 0)
	extraFrags := make(map[uint64]([]byte))
	recvLast := false

	for {
		select {
		case <-quit:
			//Trace.Println("quitting upon quit signal")
			didQuit <- true
			return
		case frag := <-in:
			//Trace.Println("got a fragment packet. len:", len(frag))
			hdr, p := slicePacket(frag)
			//offset := 8 * (uint64(hdr[6]&0x1f)<<8 + uint64(hdr[7]))
			//fmt.Println("RECEIVED FRAG")
			//fmt.Println("Offset:", offset)
			//fmt.Println(len(payload))

			// add to map
			offset := 8 * (uint64(hdr[6]&0x1F)<<8 + uint64(hdr[7]))
			//Trace.Println("Offset:", offset)
			extraFrags[offset] = p

			// check more fragments flag
			if (hdr[6]>>5)&0x01 == 0 {
				recvLast = true
			}

			// add to payload
			//Trace.Println("Begin to add to the payload")
			for {
				if storedFrag, found := extraFrags[uint64(len(payload))]; found {
					//Trace.Println("New Payload Len: ", len(payload))
					delete(extraFrags, uint64(len(payload)))
					payload = append(payload, storedFrag...)
				} else {
					break
				}
			}
			//Trace.Println("Finished add to the payload")

			// deal with the payload
			if recvLast && len(extraFrags) == 0 {
				//Trace.Println("Done")
				// correct the header
				fullPacketHdr := hdr
				totalLen := uint16(fullPacketHdr[0]&0x0F)*4 + uint16(len(payload))
				fullPacketHdr[2] = byte(totalLen >> 8)
				fullPacketHdr[3] = byte(totalLen)
				fullPacketHdr[6] = 0
				fullPacketHdr[7] = 0

				// update the checksum
				check := calculateIPChecksum(fullPacketHdr[:20])
				fullPacketHdr[10] = byte(check >> 8)
				fullPacketHdr[11] = byte(check)

				// send the packet back into processing
				//go func() {
				select {
				case ipr.incomingPackets <- append(fullPacketHdr, payload...):
				default:
					logs.Warn.Println("Dropping defragmented packet, no space in buffer")
				}
				//fmt.Println("FINISHED")
				//}()
				//Trace.Println("Just wrote back in")
				done <- true
				return // from goroutine
			}
			//Trace.Println("Looping")
		}
	}

	// drop the packet upon timeout
	//Trace.Println(errors.New("Fragments took too long, packet dropped"))
	//return
}

func (ipr *IP_Reader) killFragmentAssembler(quit chan<- bool, didQuit <-chan bool, done <-chan bool, bufID string) {
	// sends quit to the assembler if it doesn't send done
	select {
	case <-time.After(FRAGMENT_TIMEOUT):
		//Trace.Println("Force quitting packet assembler")
		quit <- true
		<-didQuit // will block until it has been received
	case <-done:
		//Trace.Println("Recieved done msg.")
	}

	//Trace.Println("Frag Assemble Ended, finished")
	delete(ipr.fragBuf, bufID)
}

func (ipr *IP_Reader) ReadFrom() (rip, lip *ipv4tps.IPaddress, b, payload []byte, e error) {
	//fmt.Println("STARTING READ")
	b = <-ipr.incomingPackets
	//	logs.Info.Println("Read IP packet")
	//fmt.Println("RAW READ COMPLETED")
	//fmt.Println("Read Length: ", len(b))
	//fmt.Print(".")
	//fmt.Println("Full Read Data: ", b)

	hdr, p := slicePacket(b)

	// extract source IP and protocol
	rip = &ipv4tps.IPaddress{IP: hdr[12:16]}
	lip = &ipv4tps.IPaddress{IP: hdr[16:20]}
	proto := uint8(hdr[9])
	if !((bytes.Equal(ipr.ip.IP, rip.IP) || bytes.Equal(ipr.ip.IP, ipv4tps.IP_ALL)) && ipr.protocol == proto) {
		//Info.Println("Not interested in packet: dropping.")
		return ipr.ReadFrom()
	}

	// verify checksum
	if !verifyIPChecksum(hdr) {
		//Info.Println("Header checksum incorrect, packet dropped")
		return ipr.ReadFrom() // return another packet instead
	}

	packetOffset := uint16(hdr[6]&0x1f)<<8 + uint16(hdr[7])
	//fmt.Println("PACK OFF", packetOffset, "HEADER FLAGS", (hdr[6] >> 5))
	if ((hdr[6]>>5)&0x01 == 0) && (packetOffset == 0) {
		// not a fragment
		return rip, lip, b, p, nil
	} else {
		// is a fragment
		bufID := string([]byte{hdr[12], hdr[13], hdr[14], hdr[15], // the source IP
			hdr[16], hdr[17], hdr[18], hdr[19], // the destination IP
			hdr[9],         // protocol
			hdr[4], hdr[5], // identification
		})
		//Trace.Printf("rcvd a fragment-bufId: %x, len: %d\n", bufID, len(b))

		if _, ok := ipr.fragBuf[bufID]; !ok {
			// create the fragment buffer and quit
			//Trace.Printf("creating a new buffer for %x\n", bufID)
			ipr.fragBuf[bufID] = make(chan []byte, FRAGMENT_ASSEMBLER_BUFFER_SIZE)
			quit := make(chan bool, 1)
			done := make(chan bool, 1)
			didQuit := make(chan bool, 1)

			// create the packet assembler in a goroutine to allow the program to continue
			go ipr.fragmentAssembler(ipr.fragBuf[bufID], quit, didQuit, done)
			go ipr.killFragmentAssembler(quit, didQuit, done, bufID)
		}

		// send the packet to the assembler
		go func() { ipr.fragBuf[bufID] <- b }()

		// after dealing with the fragment, try reading again
		return ipr.ReadFrom()
	}
}

func (ipr *IP_Reader) Close() error {
	return ipr.irm.Unbind(ipr.ip, ipr.protocol)
}

/* h := &ipv4.Header{
	Version:  ipv4.Version,      // protocol version
	Len:      20,                // header length
	TOS:      0,                 // type-of-service (0 is everything normal)
	TotalLen: len(x) + 20,       // packet total length (octets)
	ID:       0,                 // identification
	Flags:    ipv4.DontFragment, // flags
	FragOff:  0,                 // fragment offset
	TTL:      8,                 // time-to-live (maximum lifespan in seconds)
	Protocol: 17,                // next protocol (17 is UDP)
	Checksum: 0,                 // checksum (apparently autocomputed)
	//Src:    net.IPv4(127, 0, 0, 1), // source address, apparently done automatically
	Dst: net.ParseIP(c.manager.ipAddress), // destination address
	//Options                         // options, extension headers
}
*/
<|endoftext|>