<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/restic/restic"
	"github.com/restic/restic/backend"
)

type CmdRestore struct{}

func init() {
	_, err := parser.AddCommand("restore",
		"restore a snapshot",
		"The restore command restores a snapshot to a directory",
		&CmdRestore{})
	if err != nil {
		panic(err)
	}
}

func (cmd CmdRestore) Usage() string {
	return "snapshot-ID TARGETDIR"
}

func (cmd CmdRestore) Execute(args []string) error {
	if len(args) != 2 {
		return fmt.Errorf("wrong number of arguments, Usage: %s", cmd.Usage())
	}

	s, err := OpenRepo()
	if err != nil {
		return err
	}

	id, err := backend.FindSnapshot(s, args[0])
	if err != nil {
		errx(1, "invalid id %q: %v", args[0], err)
	}

	target := args[1]

	// create restorer
	res, err := restic.NewRestorer(s, id)
	if err != nil {
		fmt.Fprintf(os.Stderr, "creating restorer failed: %v\n", err)
		os.Exit(2)
	}

	res.Error = func(dir string, node *restic.Node, err error) error {
		fmt.Fprintf(os.Stderr, "error for %s: %+v\n", dir, err)

		// if node.Type == "dir" {
		// 	if e, ok := err.(*os.PathError); ok {
		// 		if errn, ok := e.Err.(syscall.Errno); ok {
		// 			if errn == syscall.EEXIST {
		// 				fmt.Printf("ignoring already existing directory %s\n", dir)
		// 				return nil
		// 			}
		// 		}
		// 	}
		// }
		return err
	}

	fmt.Printf("restoring %s to %s\n", res.Snapshot(), target)

	err = res.RestoreTo(target)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>Support for pattern in restore command<commit_after>package main

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/restic/restic"
	"github.com/restic/restic/backend"
)

type CmdRestore struct{}

func init() {
	_, err := parser.AddCommand("restore",
		"restore a snapshot",
		"The restore command restores a snapshot to a directory",
		&CmdRestore{})
	if err != nil {
		panic(err)
	}
}

func (cmd CmdRestore) Usage() string {
	return "snapshot-ID TARGETDIR [PATTERN]"
}

func (cmd CmdRestore) Execute(args []string) error {
	if len(args) < 2 || len(args) > 3 {
		return fmt.Errorf("wrong number of arguments, Usage: %s", cmd.Usage())
	}

	s, err := OpenRepo()
	if err != nil {
		return err
	}

	id, err := backend.FindSnapshot(s, args[0])
	if err != nil {
		errx(1, "invalid id %q: %v", args[0], err)
	}

	target := args[1]

	// create restorer
	res, err := restic.NewRestorer(s, id)
	if err != nil {
		fmt.Fprintf(os.Stderr, "creating restorer failed: %v\n", err)
		os.Exit(2)
	}

	res.Error = func(dir string, node *restic.Node, err error) error {
		fmt.Fprintf(os.Stderr, "error for %s: %+v\n", dir, err)

		// if node.Type == "dir" {
		// 	if e, ok := err.(*os.PathError); ok {
		// 		if errn, ok := e.Err.(syscall.Errno); ok {
		// 			if errn == syscall.EEXIST {
		// 				fmt.Printf("ignoring already existing directory %s\n", dir)
		// 				return nil
		// 			}
		// 		}
		// 	}
		// }
		return err
	}

	// TODO: a filter against the full path sucks as filepath.Match doesn't match
	// directory separators on '*'. still, it's better than nothing.
	if len(args) > 2 {
		res.Filter = func(item string, dstpath string, node *restic.Node) bool {
			matched, err := filepath.Match(item, args[2])
			if err != nil {
				panic(err)
			}
			return matched
		}
	}

	fmt.Printf("restoring %s to %s\n", res.Snapshot(), target)

	err = res.RestoreTo(target)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// dhclient sets up DHCP.
//
// Synopsis:
//     dhclient [OPTIONS...]
//
// Options:
//     -timeout:  lease timeout in seconds
//     -renewals: number of DHCP renewals before exiting
//     -verbose:  verbose output
package main

import (
	"crypto/rand"
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"regexp"
	"sync"
	"time"

	"github.com/d2g/dhcp4"
	"github.com/d2g/dhcp4client"
	"github.com/u-root/u-root/pkg/dhcp6client"
	"github.com/vishvananda/netlink"
)

const (
	// slop is the slop in our lease time.
	slop          = 10 * time.Second
	linkUpAttempt = 30 * time.Second
)

var (
	ifName       = "^e.*"
	leasetimeout = flag.Int("timeout", 15, "Lease timeout in seconds")
	retry        = flag.Int("retry", -1, "Max number of attempts for DHCP clients to send requests. -1 means infinity")
	renewals     = flag.Int("renewals", -1, "Number of DHCP renewals before exiting. -1 means infinity")
	verbose      = flag.Bool("verbose", false, "Verbose output")
	ipv4         = flag.Bool("ipv4", false, "use IPV4")
	test         = flag.Bool("test", false, "Test mode")
	debug        = func(string, ...interface{}) {}
)

func ifup(ifname string) (netlink.Link, error) {
	debug("Try bringing up %v", ifname)
	start := time.Now()
	for time.Since(start) < linkUpAttempt {
		// Note that it may seem odd to keep trying the
		// LinkByName operation, by consider that a hotplug
		// device such as USB ethernet can just vanish.
		iface, err := netlink.LinkByName(ifname)
		debug("LinkByName(%v) returns (%v, %v)", ifname, iface, err)
		if err != nil {
			return nil, fmt.Errorf("cannot get interface by name %v: %v", ifname, err)
		}

		if iface.Attrs().OperState == netlink.OperUp {
			debug("Link %v is up", ifname)
			return iface, nil
		}

		if err := netlink.LinkSetUp(iface); err != nil {
			return nil, fmt.Errorf("%v: %v can't make it up: %v", ifname, iface, err)
		}
		time.Sleep(1 * time.Second)
	}
	return nil, fmt.Errorf("Link %v still down after %d seconds", ifname, linkUpAttempt)
}

func dhclient4(iface netlink.Link, numRenewals int, timeout time.Duration, retry int) error {
	mac := iface.Attrs().HardwareAddr
	conn, err := dhcp4client.NewPacketSock(iface.Attrs().Index)
	if err != nil {
		return fmt.Errorf("client conection generation: %v", err)
	}

	client, err := dhcp4client.New(dhcp4client.HardwareAddr(mac), dhcp4client.Connection(conn), dhcp4client.Timeout(timeout))
	if err != nil {
		return fmt.Errorf("error: %v", err)
	}

	var packet dhcp4.Packet
	needsRequest := true
	for i := 0; numRenewals < 0 || i < numRenewals+1; i++ {
		debug("Start getting or renewing DHCPv4 lease")

		var success bool
		for i := 0; i < retry || retry < 0; i++ {
			if i > 0 {
				if needsRequest {
					debug("Resending DHCPv4 request...\n")
				} else {
					debug("Resending DHCPv4 renewal")
				}
			}

			if needsRequest {
				success, packet, err = client.Request()
			} else {
				success, packet, err = client.Renew(packet)
			}
			if err != nil {
				if err0, ok := err.(net.Error); ok && err0.Timeout() {
					log.Printf("%s: timeout contacting DHCP server", mac)
				} else {
					log.Printf("%s: error: %v", mac, err)
				}
			} else {
				// Client needs renew after no matter what state it is now.
				needsRequest = false
				break
			}
		}

		debug("Success on %s: %v\n", mac, success)
		debug("Packet: %v\n", packet)
		debug("Lease is %v seconds\n", packet.Secs())

		if !success {
			return fmt.Errorf("%s: we didn't sucessfully get a DHCP lease", mac)
		}
		debug("IP Received: %v\n", packet.YIAddr().String())

		// We got here because we got a good packet.
		o := packet.ParseOptions()

		netmask, ok := o[dhcp4.OptionSubnetMask]
		if ok {
			debug("OptionSubnetMask is %v\n", netmask)
		} else {
			// If they did not offer a subnet mask, we
			// choose the most restrictive option, namely,
			// our IP address.  This could happen on,
			// e.g., a point to point link.
			netmask = packet.YIAddr()
			debug("No OptionSubnetMask; default to %v\n", netmask)
		}

		dst := &netlink.Addr{IPNet: &net.IPNet{IP: packet.YIAddr(), Mask: netmask}, Label: ""}
		// Add the address to the iface.
		if *test == false {
			if err := netlink.AddrReplace(iface, dst); err != nil {
				if os.IsExist(err) {
					return fmt.Errorf("add/replace %v to %v: %v", dst, iface, err)
				}
			}

			if gwData, ok := o[dhcp4.OptionRouter]; ok {
				debug("router %v", gwData)
				routerName := net.IP(gwData).String()
				debug("routerName %v", routerName)
				r := &netlink.Route{
					Dst:       &net.IPNet{IP: packet.GIAddr()},
					LinkIndex: iface.Attrs().Index,
					Gw:        net.IP(gwData),
				}

				if err := netlink.RouteReplace(r); err != nil {
					return fmt.Errorf("%s: add %s: %v", iface.Attrs().Name, r.String(), routerName)
				}
			}
		}
		if binary.BigEndian.Uint16(packet.Secs()) == 0 {
			debug("%v: server returned infinite lease.", iface.Attrs().Name)
			break
		}

		// We can not assume the server will give us any grace time. So
		// sleep for just a tiny bit less than the minimum.
		time.Sleep(timeout - slop)
	}
	return nil
}

// dhcp6 support in go is hard to find. This function represents our best current
// guess based on reading and testing.
func dhclient6(iface netlink.Link, numRenewals int, timeout time.Duration, retry int) error {
	conn, err := dhcp6client.NewPacketSock(iface.Attrs().Index)
	if err != nil {
		return fmt.Errorf("client conection generation: %v", err)
	}
	client := dhcp6client.New(iface.Attrs().HardwareAddr, conn, timeout, retry)

	for i := 0; numRenewals < 0 || i < numRenewals+1; i++ {
		debug("Start getting or renewing DHCPv6 lease")
		iaAddrs, packet, err := client.Solicit()
		if err != nil {
			return fmt.Errorf("error: %v", err)
		}
		debug("Packet: %+v\n\n", packet)
		debug("IAAddrs: %v\n", iaAddrs)

		if *test == false {
			dst := &netlink.Addr{
				IPNet:       &net.IPNet{IP: iaAddrs[0].IP},
				PreferedLft: int(iaAddrs[0].PreferredLifetime.Seconds()),
				ValidLft:    int(iaAddrs[0].ValidLifetime.Seconds()),
				Label:       "",
			}

			if err := netlink.AddrReplace(iface, dst); err != nil {
				if os.IsExist(err) {
					return fmt.Errorf("add/replace %v to %v: %v", dst, iface, err)
				}
			}
		}

		time.Sleep(timeout - slop)
	}
	return nil
}

func main() {
	flag.Parse()
	if *verbose {
		debug = log.Printf
	}

	// if we boot quickly enough, the random number generator
	// may not be ready, and the dhcp package panics in that case.
	// Worse, /dev/urandom, which the Go package falls back to,
	// might not be there. Still worse, the Go package is "sticky"
	// in that once it decides to use /dev/urandom, it won't go back,
	// even if the system call would subsequently work.
	// You're screwed. Exit.
	// Wouldn't it be nice if we could just do the blocking system
	// call? But that comes with its own giant set of headaches.
	// Maybe we'll end up in a loop, sleeping, and just running
	// ourselves.
	if n, err := rand.Read([]byte{0}); err != nil || n != 1 {
		log.Fatalf("We're sorry, the random number generator is not up. Please file a ticket")
	}

	if len(flag.Args()) > 1 {
		log.Fatalf("only one re")
	}

	if len(flag.Args()) > 0 {
		ifName = flag.Args()[0]
	}

	ifRE := regexp.MustCompilePOSIX(ifName)

	ifnames, err := netlink.LinkList()
	if err != nil {
		log.Fatalf("Can't get list of link names: %v", err)
	}

	timeout := time.Duration(*leasetimeout) * time.Second
	// if timeout is < slop, it's too short.
	if timeout < slop {
		timeout = 2 * slop
		log.Printf("increased lease timeout to %s", timeout)
	}

	var wg sync.WaitGroup
	done := make(chan error)
	for _, i := range ifnames {
		if !ifRE.MatchString(i.Attrs().Name) {
			continue
		}
		wg.Add(1)
		go func(ifname string) {
			defer wg.Done()
			iface, err := ifup(ifname)
			if err != nil {
				done <- err
				return
			}
			if *ipv4 {
				done <- dhclient4(iface, *renewals, timeout, *retry)
			} else {
				done <- dhclient6(iface, *renewals, timeout, *retry)
			}
			debug("Done dhclient for %v", ifname)
		}(i.Attrs().Name)
	}

	go func() {
		wg.Wait()
		close(done)
	}()
	// Wait for all goroutines to finish.
	var nif int
	for err := range done {
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
		}
		nif++
	}

	if nif == 0 {
		fmt.Printf("No interfaces match %v\n", ifName)
	}
}
<commit_msg>dhclient: enable both ipv4 and ipv6 by default<commit_after>// Copyright 2017 the u-root Authors. All rights reserved
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// dhclient sets up DHCP.
//
// Synopsis:
//     dhclient [OPTIONS...]
//
// Options:
//     -timeout:  lease timeout in seconds
//     -renewals: number of DHCP renewals before exiting
//     -verbose:  verbose output
package main

import (
	"crypto/rand"
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"regexp"
	"sync"
	"time"

	"github.com/d2g/dhcp4"
	"github.com/d2g/dhcp4client"
	"github.com/u-root/u-root/pkg/dhcp6client"
	"github.com/vishvananda/netlink"
)

const (
	// slop is the slop in our lease time.
	slop          = 10 * time.Second
	linkUpAttempt = 30 * time.Second
)

var (
	ifName       = "^e.*"
	leasetimeout = flag.Int("timeout", 15, "Lease timeout in seconds")
	retry        = flag.Int("retry", -1, "Max number of attempts for DHCP clients to send requests. -1 means infinity")
	renewals     = flag.Int("renewals", -1, "Number of DHCP renewals before exiting. -1 means infinity")
	verbose      = flag.Bool("verbose", false, "Verbose output")
	ipv4         = flag.Bool("ipv4", true, "use IPV4")
	ipv6         = flag.Bool("ipv6", true, "use IPV6")
	test         = flag.Bool("test", false, "Test mode")
	debug        = func(string, ...interface{}) {}
)

func ifup(ifname string) (netlink.Link, error) {
	debug("Try bringing up %v", ifname)
	start := time.Now()
	for time.Since(start) < linkUpAttempt {
		// Note that it may seem odd to keep trying the
		// LinkByName operation, by consider that a hotplug
		// device such as USB ethernet can just vanish.
		iface, err := netlink.LinkByName(ifname)
		debug("LinkByName(%v) returns (%v, %v)", ifname, iface, err)
		if err != nil {
			return nil, fmt.Errorf("cannot get interface by name %v: %v", ifname, err)
		}

		if iface.Attrs().OperState == netlink.OperUp {
			debug("Link %v is up", ifname)
			return iface, nil
		}

		if err := netlink.LinkSetUp(iface); err != nil {
			return nil, fmt.Errorf("%v: %v can't make it up: %v", ifname, iface, err)
		}
		time.Sleep(1 * time.Second)
	}
	return nil, fmt.Errorf("Link %v still down after %d seconds", ifname, linkUpAttempt)
}

func dhclient4(iface netlink.Link, numRenewals int, timeout time.Duration, retry int) error {
	mac := iface.Attrs().HardwareAddr
	conn, err := dhcp4client.NewPacketSock(iface.Attrs().Index)
	if err != nil {
		return fmt.Errorf("client conection generation: %v", err)
	}

	client, err := dhcp4client.New(dhcp4client.HardwareAddr(mac), dhcp4client.Connection(conn), dhcp4client.Timeout(timeout))
	if err != nil {
		return fmt.Errorf("error: %v", err)
	}

	var packet dhcp4.Packet
	needsRequest := true
	for i := 0; numRenewals < 0 || i < numRenewals+1; i++ {
		debug("Start getting or renewing DHCPv4 lease")

		var success bool
		for i := 0; i < retry || retry < 0; i++ {
			if i > 0 {
				if needsRequest {
					debug("Resending DHCPv4 request...\n")
				} else {
					debug("Resending DHCPv4 renewal")
				}
			}

			if needsRequest {
				success, packet, err = client.Request()
			} else {
				success, packet, err = client.Renew(packet)
			}
			if err != nil {
				if err0, ok := err.(net.Error); ok && err0.Timeout() {
					log.Printf("%s: timeout contacting DHCP server", mac)
				} else {
					log.Printf("%s: error: %v", mac, err)
				}
			} else {
				// Client needs renew after no matter what state it is now.
				needsRequest = false
				break
			}
		}

		debug("Success on %s: %v\n", mac, success)
		debug("Packet: %v\n", packet)
		debug("Lease is %v seconds\n", packet.Secs())

		if !success {
			return fmt.Errorf("%s: we didn't sucessfully get a DHCP lease", mac)
		}
		debug("IP Received: %v\n", packet.YIAddr().String())

		// We got here because we got a good packet.
		o := packet.ParseOptions()

		netmask, ok := o[dhcp4.OptionSubnetMask]
		if ok {
			debug("OptionSubnetMask is %v\n", netmask)
		} else {
			// If they did not offer a subnet mask, we
			// choose the most restrictive option, namely,
			// our IP address.  This could happen on,
			// e.g., a point to point link.
			netmask = packet.YIAddr()
			debug("No OptionSubnetMask; default to %v\n", netmask)
		}

		dst := &netlink.Addr{IPNet: &net.IPNet{IP: packet.YIAddr(), Mask: netmask}, Label: ""}
		// Add the address to the iface.
		if *test == false {
			if err := netlink.AddrReplace(iface, dst); err != nil {
				if os.IsExist(err) {
					return fmt.Errorf("add/replace %v to %v: %v", dst, iface, err)
				}
			}

			if gwData, ok := o[dhcp4.OptionRouter]; ok {
				debug("router %v", gwData)
				routerName := net.IP(gwData).String()
				debug("routerName %v", routerName)
				r := &netlink.Route{
					Dst:       &net.IPNet{IP: packet.GIAddr()},
					LinkIndex: iface.Attrs().Index,
					Gw:        net.IP(gwData),
				}

				if err := netlink.RouteReplace(r); err != nil {
					return fmt.Errorf("%s: add %s: %v", iface.Attrs().Name, r.String(), routerName)
				}
			}
		}
		if binary.BigEndian.Uint16(packet.Secs()) == 0 {
			debug("%v: server returned infinite lease.", iface.Attrs().Name)
			break
		}

		// We can not assume the server will give us any grace time. So
		// sleep for just a tiny bit less than the minimum.
		time.Sleep(timeout - slop)
	}
	return nil
}

// dhcp6 support in go is hard to find. This function represents our best current
// guess based on reading and testing.
func dhclient6(iface netlink.Link, numRenewals int, timeout time.Duration, retry int) error {
	conn, err := dhcp6client.NewPacketSock(iface.Attrs().Index)
	if err != nil {
		return fmt.Errorf("client conection generation: %v", err)
	}
	client := dhcp6client.New(iface.Attrs().HardwareAddr, conn, timeout, retry)

	for i := 0; numRenewals < 0 || i < numRenewals+1; i++ {
		debug("Start getting or renewing DHCPv6 lease")
		iaAddrs, packet, err := client.Solicit()
		if err != nil {
			return fmt.Errorf("error: %v", err)
		}
		debug("Packet: %+v\n\n", packet)
		debug("IAAddrs: %v\n", iaAddrs)

		if *test == false {
			dst := &netlink.Addr{
				IPNet:       &net.IPNet{IP: iaAddrs[0].IP},
				PreferedLft: int(iaAddrs[0].PreferredLifetime.Seconds()),
				ValidLft:    int(iaAddrs[0].ValidLifetime.Seconds()),
				Label:       "",
			}

			if err := netlink.AddrReplace(iface, dst); err != nil {
				if os.IsExist(err) {
					return fmt.Errorf("add/replace %v to %v: %v", dst, iface, err)
				}
			}
		}

		time.Sleep(timeout - slop)
	}
	return nil
}

func main() {
	flag.Parse()
	if *verbose {
		debug = log.Printf
	}

	// if we boot quickly enough, the random number generator
	// may not be ready, and the dhcp package panics in that case.
	// Worse, /dev/urandom, which the Go package falls back to,
	// might not be there. Still worse, the Go package is "sticky"
	// in that once it decides to use /dev/urandom, it won't go back,
	// even if the system call would subsequently work.
	// You're screwed. Exit.
	// Wouldn't it be nice if we could just do the blocking system
	// call? But that comes with its own giant set of headaches.
	// Maybe we'll end up in a loop, sleeping, and just running
	// ourselves.
	if n, err := rand.Read([]byte{0}); err != nil || n != 1 {
		log.Fatalf("We're sorry, the random number generator is not up. Please file a ticket")
	}

	if len(flag.Args()) > 1 {
		log.Fatalf("only one re")
	}

	if len(flag.Args()) > 0 {
		ifName = flag.Args()[0]
	}

	ifRE := regexp.MustCompilePOSIX(ifName)

	ifnames, err := netlink.LinkList()
	if err != nil {
		log.Fatalf("Can't get list of link names: %v", err)
	}

	timeout := time.Duration(*leasetimeout) * time.Second
	// if timeout is < slop, it's too short.
	if timeout < slop {
		timeout = 2 * slop
		log.Printf("increased lease timeout to %s", timeout)
	}

	var wg sync.WaitGroup
	done := make(chan error)
	for _, i := range ifnames {
		if !ifRE.MatchString(i.Attrs().Name) {
			continue
		}
		wg.Add(1)
		go func(ifname string) {
			defer wg.Done()
			iface, err := ifup(ifname)
			if err != nil {
				done <- err
				return
			}
			if *ipv4 {
				wg.Add(1)
				done <- dhclient4(iface, *renewals, timeout, *retry)
				wg.Done()
			}
			if *ipv6 {
				wg.Add(1)
				done <- dhclient6(iface, *renewals, timeout, *retry)
				wg.Done()
			}
			debug("Done dhclient for %v", ifname)
		}(i.Attrs().Name)
	}

	go func() {
		wg.Wait()
		close(done)
	}()
	// Wait for all goroutines to finish.
	var nif int
	for err := range done {
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
		}
		nif++
	}

	if nif == 0 {
		fmt.Printf("No interfaces match %v\n", ifName)
	}
}
<|endoftext|>


<commit_before>package shells

import (
	"bufio"
	"bytes"
	"fmt"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers"
	"io"
	"path/filepath"
	"runtime"
	"strings"
)

type BashShell struct {
	AbstractShell
}

func (b *BashShell) GetName() string {
	return "bash"
}

func (b *BashShell) writeCloneCmd(w io.Writer, build *common.Build, projectDir string) {
	io.WriteString(w, "echo Cloning repository...\n")
	io.WriteString(w, fmt.Sprintf("rm -rf %s\n", projectDir))
	io.WriteString(w, fmt.Sprintf("mkdir -p %s\n", projectDir))
	io.WriteString(w, fmt.Sprintf("git clone %s %s\n", helpers.ShellEscape(build.RepoURL), projectDir))
	io.WriteString(w, fmt.Sprintf("cd %s\n", projectDir))
}

func (b *BashShell) writeFetchCmd(w io.Writer, build *common.Build, projectDir string, gitDir string) {
	io.WriteString(w, fmt.Sprintf("if [[ -d %s ]]; then\n", gitDir))
	io.WriteString(w, "echo Fetching changes...\n")
	io.WriteString(w, fmt.Sprintf("cd %s\n", projectDir))
	io.WriteString(w, fmt.Sprintf("git clean -fdx\n"))
	io.WriteString(w, fmt.Sprintf("git reset --hard > /dev/null\n"))
	io.WriteString(w, fmt.Sprintf("git remote set-url origin %s\n", helpers.ShellEscape(build.RepoURL)))
	io.WriteString(w, fmt.Sprintf("git fetch origin\n"))
	io.WriteString(w, fmt.Sprintf("else\n"))
	b.writeCloneCmd(w, build, projectDir)
	io.WriteString(w, fmt.Sprintf("fi\n"))
}

func (b *BashShell) writeCheckoutCmd(w io.Writer, build *common.Build) {
	io.WriteString(w, fmt.Sprintf("echo Checking out %s as %s...\n", build.Sha[0:8], build.RefName))
	io.WriteString(w, fmt.Sprintf("git checkout -qf %s\n", build.Sha))
}

func (b *BashShell) GenerateScript(info common.ShellScriptInfo) (*common.ShellScript, error) {
	var buffer bytes.Buffer
	w := bufio.NewWriter(&buffer)

	build := info.Build
	projectDir := build.FullProjectDir()
	projectDir = helpers.ToSlash(projectDir)
	gitDir := filepath.Join(projectDir, ".git")

	projectScript := helpers.ShellEscape(build.FullProjectDir() + ".sh")

	io.WriteString(w, "#!/usr/bin/env bash\n")
	io.WriteString(w, "\n")
	io.WriteString(w, "# save script that is read from to file and execute script file on remote server\n")
	io.WriteString(w, fmt.Sprintf("mkdir -p %s\n", helpers.ShellEscape(projectDir)))
	io.WriteString(w, fmt.Sprintf("cat > %s && exec $0 %s; exit 1\n", projectScript, projectScript))
	io.WriteString(w, "\n")
	if len(build.Hostname) != 0 {
		io.WriteString(w, fmt.Sprintf("echo Running on $(hostname) via %s...\n", helpers.ShellEscape(build.Hostname)))
	} else {
		io.WriteString(w, "echo Running on $(hostname)...\n")
	}
	io.WriteString(w, "\n")
	io.WriteString(w, "set -eo pipefail\n")

	io.WriteString(w, "\n")
	if build.AllowGitFetch {
		b.writeFetchCmd(w, build, helpers.ShellEscape(projectDir), helpers.ShellEscape(gitDir))
	} else {
		b.writeCloneCmd(w, build, helpers.ShellEscape(projectDir))
	}

	b.writeCheckoutCmd(w, build)
	io.WriteString(w, "\n")
	if !helpers.BoolOrDefault(build.Runner.DisableVerbose, false) {
		io.WriteString(w, "set -v\n")
		io.WriteString(w, "\n")
	}

	commands := build.Commands
	commands = strings.Replace(commands, "\r\n", "\n", -1)
	io.WriteString(w, commands)
	io.WriteString(w, "\n")

	w.Flush()

	script := common.ShellScript{
		Environment: b.GetVariables(build, projectDir),
		Script:      buffer.String(),
	}

	// su
	if info.User != nil {
		script.Command = "su"
		if info.Type == common.LoginShell {
			script.Arguments = []string{"--shell", "/bin/bash", "--login", *info.User}
		} else {
			script.Arguments = []string{"--shell", "/bin/bash", *info.User}
		}
	} else {
		script.Command = "bash"
		if info.Type == common.LoginShell {
			script.Arguments = []string{"--login"}
		}
	}

	return &script, nil
}

func (b *BashShell) IsDefault() bool {
	return runtime.GOOS != "windows"
}

func init() {
	common.RegisterShell(&BashShell{})
}
<commit_msg>Fix bash service execution<commit_after>package shells

import (
	"bufio"
	"bytes"
	"fmt"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/common"
	"gitlab.com/gitlab-org/gitlab-ci-multi-runner/helpers"
	"io"
	"path/filepath"
	"runtime"
	"strings"
)

type BashShell struct {
	AbstractShell
}

func (b *BashShell) GetName() string {
	return "bash"
}

func (b *BashShell) writeCloneCmd(w io.Writer, build *common.Build, projectDir string) {
	io.WriteString(w, "echo Cloning repository...\n")
	io.WriteString(w, fmt.Sprintf("rm -rf %s\n", projectDir))
	io.WriteString(w, fmt.Sprintf("mkdir -p %s\n", projectDir))
	io.WriteString(w, fmt.Sprintf("git clone %s %s\n", helpers.ShellEscape(build.RepoURL), projectDir))
	io.WriteString(w, fmt.Sprintf("cd %s\n", projectDir))
}

func (b *BashShell) writeFetchCmd(w io.Writer, build *common.Build, projectDir string, gitDir string) {
	io.WriteString(w, fmt.Sprintf("if [[ -d %s ]]; then\n", gitDir))
	io.WriteString(w, "echo Fetching changes...\n")
	io.WriteString(w, fmt.Sprintf("cd %s\n", projectDir))
	io.WriteString(w, fmt.Sprintf("git clean -fdx\n"))
	io.WriteString(w, fmt.Sprintf("git reset --hard > /dev/null\n"))
	io.WriteString(w, fmt.Sprintf("git remote set-url origin %s\n", helpers.ShellEscape(build.RepoURL)))
	io.WriteString(w, fmt.Sprintf("git fetch origin\n"))
	io.WriteString(w, fmt.Sprintf("else\n"))
	b.writeCloneCmd(w, build, projectDir)
	io.WriteString(w, fmt.Sprintf("fi\n"))
}

func (b *BashShell) writeCheckoutCmd(w io.Writer, build *common.Build) {
	io.WriteString(w, fmt.Sprintf("echo Checking out %s as %s...\n", build.Sha[0:8], build.RefName))
	io.WriteString(w, fmt.Sprintf("git checkout -qf %s\n", build.Sha))
}

func (b *BashShell) GenerateScript(info common.ShellScriptInfo) (*common.ShellScript, error) {
	var buffer bytes.Buffer
	w := bufio.NewWriter(&buffer)

	build := info.Build
	projectDir := build.FullProjectDir()
	projectDir = helpers.ToSlash(projectDir)
	gitDir := filepath.Join(projectDir, ".git")

	projectScript := helpers.ShellEscape(build.FullProjectDir() + ".sh")

	io.WriteString(w, "#!/usr/bin/env bash\n")
	io.WriteString(w, "\n")
	io.WriteString(w, "# save script that is read from to file and execute script file on remote server\n")
	io.WriteString(w, fmt.Sprintf("mkdir -p %s\n", helpers.ShellEscape(projectDir)))
	io.WriteString(w, fmt.Sprintf("cat > %s && exec /usr/bin/env bash %s; exit 1\n", projectScript, projectScript))
	io.WriteString(w, "\n")
	if len(build.Hostname) != 0 {
		io.WriteString(w, fmt.Sprintf("echo Running on $(hostname) via %s...\n", helpers.ShellEscape(build.Hostname)))
	} else {
		io.WriteString(w, "echo Running on $(hostname)...\n")
	}
	io.WriteString(w, "\n")
	io.WriteString(w, "set -eo pipefail\n")

	io.WriteString(w, "\n")
	if build.AllowGitFetch {
		b.writeFetchCmd(w, build, helpers.ShellEscape(projectDir), helpers.ShellEscape(gitDir))
	} else {
		b.writeCloneCmd(w, build, helpers.ShellEscape(projectDir))
	}

	b.writeCheckoutCmd(w, build)
	io.WriteString(w, "\n")
	if !helpers.BoolOrDefault(build.Runner.DisableVerbose, false) {
		io.WriteString(w, "set -v\n")
		io.WriteString(w, "\n")
	}

	commands := build.Commands
	commands = strings.Replace(commands, "\r\n", "\n", -1)
	io.WriteString(w, commands)
	io.WriteString(w, "\n")

	w.Flush()

	script := common.ShellScript{
		Environment: b.GetVariables(build, projectDir),
		Script:      buffer.String(),
	}

	// su
	if info.User != nil {
		script.Command = "su"
		if info.Type == common.LoginShell {
			script.Arguments = []string{"--shell", "/bin/bash", "--login", *info.User}
		} else {
			script.Arguments = []string{"--shell", "/bin/bash", *info.User}
		}
	} else {
		script.Command = "bash"
		if info.Type == common.LoginShell {
			script.Arguments = []string{"--login"}
		}
	}

	return &script, nil
}

func (b *BashShell) IsDefault() bool {
	return runtime.GOOS != "windows"
}

func init() {
	common.RegisterShell(&BashShell{})
}
<|endoftext|>


<commit_before>package sia

import (
	"errors"
	"net"
	"reflect"
	"strconv"
	"time"
)

const (
	timeout   = time.Second * 5
	maxMsgLen = 1 << 16
)

// A NetAddress contains the information needed to contact a peer over TCP.
type NetAddress struct {
	Host string
	Port uint16
}

// String returns the NetAddress as a string, concatentating the hostname and
// port number.
func (na *NetAddress) String() string {
	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.Port)))
}

// Call establishes a TCP connection to the NetAddress, calls the provided
// function on it, and closes the connection.
func (na *NetAddress) Call(fn func(net.Conn) error) error {
	conn, err := net.DialTimeout("tcp", na.String(), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	return fn(conn)
}

// SendVal sends a value to a NetAddress. It prefixes the encoded data with a
// header, comprising a message type and message length.
func (na *NetAddress) SendVal(t byte, val interface{}) error {
	conn, err := net.DialTimeout("tcp", na.String(), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	encVal := Marshal(val)
	encLen := EncUint64(uint64(len(encVal)))
	_, err = conn.Write(append([]byte{t},
		append(encLen[:4], encVal...)...))
	return err
}

// TBD
var BootstrapPeers = []NetAddress{}

// A TCPServer sends and receives messages. It also maintains an address book
// of peers to broadcast to and make requests of.
type TCPServer struct {
	net.Listener
	myAddr      NetAddress
	addressbook map[NetAddress]struct{}
	handlerMap  map[byte]func(net.Conn, []byte) error
}

// Broadcast calls the specified function on each peer in the address book.
func (tcps *TCPServer) Broadcast(fn func(net.Conn) error) {
	for addr := range tcps.addressbook {
		addr.Call(fn)
	}
}

// Register registers a message type with a message handler. The existing
// handler for that type will be overwritten.
func (tcps *TCPServer) Register(t byte, fn func(net.Conn, []byte) error) {
	tcps.handlerMap[t] = fn
}

// RegisterSimple is for simple handlers. A simple handler decodes the message
// data and passes it to fn. fn must have the type signature:
//   func(Type) error
// i.e. a 1-adic function that returns an error.
func (tcps *TCPServer) RegisterSimple(t byte, fn interface{}) error {
	// if fn not correct type, panic
	val, typ := reflect.ValueOf(fn), reflect.TypeOf(fn)
	if typ.Kind() != reflect.Func || typ.NumIn() != 1 ||
		typ.NumOut() != 1 || typ.Out(0) != reflect.TypeOf((*error)(nil)).Elem() {
		return errors.New("registered function has wrong type signature")
	}

	// create function:
	sfn := func(_ net.Conn, b []byte) error {
		v := reflect.New(typ.In(0)).Elem()
		Unmarshal(b, v.Interface())
		err := val.Call([]reflect.Value{v})[0]
		return err.Interface().(error)
	}

	tcps.Register(t, sfn)
	return nil
}

// NewTCPServer creates a TCPServer that listens on the specified port.
func NewTCPServer(port uint16) (tcps *TCPServer, err error) {
	tcpServ, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		return
	}
	tcps = &TCPServer{
		Listener:    tcpServ,
		addressbook: make(map[NetAddress]struct{}),
		// default handlers
		handlerMap: map[byte]func(net.Conn, []byte) error{
			'H': sendHostname,
			'P': tcps.sharePeers,
		},
	}

	// spawn listener
	go tcps.listen()
	return
}

// listen runs in the background, accepting incoming connections and serving
// them. listen will return after TCPServer.Close() is called, because the
// Accept() call will fail.
func (tcps *TCPServer) listen() {
	for {
		conn, err := tcps.Accept()
		if err != nil {
			return
		}
		// it is the handler's responsibility to close the connection
		go tcps.handleConn(conn)
	}
}

// handleConn reads header data from a connection, unmarshals the data
// structures it contains, and routes the data to other functions for
// processing.
// TODO: set deadlines?
func (tcps *TCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	var (
		msgHead []byte = make([]byte, 5)
		msgData []byte // length determined by msgHead
	)
	if n, err := conn.Read(msgHead); err != nil || n != 5 {
		// TODO: log error
		return
	}
	msgLen := DecUint64(msgHead[1:])
	if msgLen > maxMsgLen {
		// TODO: log error
		return
	}
	msgData = make([]byte, msgLen)
	if n, err := conn.Read(msgData); err != nil || uint64(n) != msgLen {
		// TODO: log error
		return
	}

	// call registered handler for this message type
	if fn, ok := tcps.handlerMap[msgHead[0]]; ok {
		fn(conn, msgData)
		// TODO: log error
	}
	return
}

// sendHostname replies to the send with the sender's external IP.
func sendHostname(conn net.Conn, _ []byte) error {
	_, err := conn.Write([]byte(conn.RemoteAddr().String()))
	return err
}

// sharePeers transmits at most 'num' peers over the connection.
// TODO: choose random peers?
func (tcps *TCPServer) sharePeers(conn net.Conn, msgData []byte) error {
	if len(msgData) != 1 {
		return errors.New("invalid number of peers")
	}
	num := msgData[0]
	var addrs []NetAddress
	for addr := range tcps.addressbook {
		if num == 0 {
			break
		}
		addrs = append(addrs, addr)
		num--
	}
	_, err := conn.Write(Marshal(addrs))
	return err
}

// Ping returns whether a NetAddress is reachable. It accomplishes this by
// initiating a TCP connection and immediately closes it. This is pretty
// unsophisticated. I'll add a Pong later.
func (tcps *TCPServer) Ping(addr NetAddress) bool {
	conn, err := net.DialTimeout("tcp", addr.String(), timeout)
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

// learnHostname learns the external IP of the TCPServer.
func (tcps *TCPServer) learnHostname(conn net.Conn) (err error) {
	// send hostname request
	if _, err = conn.Write([]byte{'H', 0, 0, 0, 0}); err != nil {
		return
	}
	// read response
	buf := make([]byte, 128)
	n, err := conn.Read(buf)
	if err != nil {
		return
	}
	// TODO: try to ping ourselves?
	host, portStr, err := net.SplitHostPort(string(buf[:n]))
	if err != nil {
		return
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return
	}
	tcps.myAddr = NetAddress{host, uint16(port)}
	return
}

// requestPeers queries a peer for additional peers, and adds any new peers to
// the address book.
func (tcps *TCPServer) requestPeers(conn net.Conn) (err error) {
	// request 10 peers
	if _, err = conn.Write([]byte{'P', 1, 0, 0, 0, 10}); err != nil {
		return
	}
	// read response
	buf := make([]byte, 1024)
	n, err := conn.Read(buf)
	if err != nil {
		return
	}
	var addrs []NetAddress
	if err = Unmarshal(buf[:n], &addrs); err != nil {
		return
	}
	// add peers
	// TODO: make sure we don't add ourself
	for _, addr := range addrs {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}
	return
}

// Bootstrap discovers the external IP of the TCPServer, requests peers from
// the initial peer list, and announces itself to those peers.
func (tcps *TCPServer) Bootstrap() (err error) {
	// populate initial peer list
	for _, addr := range BootstrapPeers {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}

	// learn hostname
	for addr := range tcps.addressbook {
		if addr.Call(tcps.learnHostname) == nil {
			break
		}
	}
	// request peers
	// TODO: maybe iterate until we have enough new peers?
	for addr := range tcps.addressbook {
		addr.Call(tcps.requestPeers)
	}
	// TODO: announce ourselves to new peers
	return
}
<commit_msg>fix bugs revealed by testing<commit_after>package sia

import (
	"errors"
	"net"
	"reflect"
	"strconv"
	"time"
)

const (
	timeout   = time.Second * 5
	maxMsgLen = 1 << 16
)

// A NetAddress contains the information needed to contact a peer over TCP.
type NetAddress struct {
	Host string
	Port uint16
}

// String returns the NetAddress as a string, concatentating the hostname and
// port number.
func (na *NetAddress) String() string {
	return net.JoinHostPort(na.Host, strconv.Itoa(int(na.Port)))
}

// Call establishes a TCP connection to the NetAddress, calls the provided
// function on it, and closes the connection.
func (na *NetAddress) Call(fn func(net.Conn) error) error {
	conn, err := net.DialTimeout("tcp", na.String(), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	return fn(conn)
}

// SendVal sends a value to a NetAddress. It prefixes the encoded data with a
// header, comprising a message type and message length.
func (na *NetAddress) SendVal(t byte, val interface{}) error {
	conn, err := net.DialTimeout("tcp", na.String(), timeout)
	if err != nil {
		return err
	}
	defer conn.Close()
	encVal := Marshal(val)
	encLen := EncUint64(uint64(len(encVal)))
	_, err = conn.Write(append([]byte{t},
		append(encLen[:4], encVal...)...))
	return err
}

// TBD
var BootstrapPeers = []NetAddress{}

// A TCPServer sends and receives messages. It also maintains an address book
// of peers to broadcast to and make requests of.
type TCPServer struct {
	net.Listener
	myAddr      NetAddress
	addressbook map[NetAddress]struct{}
	handlerMap  map[byte]func(net.Conn, []byte) error
}

// Broadcast calls the specified function on each peer in the address book.
func (tcps *TCPServer) Broadcast(fn func(net.Conn) error) {
	for addr := range tcps.addressbook {
		addr.Call(fn)
	}
}

// Register registers a message type with a message handler. The existing
// handler for that type will be overwritten.
func (tcps *TCPServer) Register(t byte, fn func(net.Conn, []byte) error) {
	tcps.handlerMap[t] = fn
}

// RegisterSimple is for simple handlers. A simple handler decodes the message
// data and passes it to fn. fn must have the type signature:
//   func(Type) error
// i.e. a 1-adic function that returns an error.
func (tcps *TCPServer) RegisterSimple(t byte, fn interface{}) error {
	// if fn not correct type, panic
	val, typ := reflect.ValueOf(fn), reflect.TypeOf(fn)
	if typ.Kind() != reflect.Func || typ.NumIn() != 1 ||
		typ.NumOut() != 1 || typ.Out(0) != reflect.TypeOf((*error)(nil)).Elem() {
		return errors.New("registered function has wrong type signature")
	}

	// create function:
	sfn := func(_ net.Conn, b []byte) error {
		v := reflect.New(typ.In(0)).Elem()
		Unmarshal(b, v.Interface())
		if err := val.Call([]reflect.Value{v})[0].Interface(); err != nil {
			return err.(error)
		}
		return nil
	}

	tcps.Register(t, sfn)
	return nil
}

// NewTCPServer creates a TCPServer that listens on the specified port.
func NewTCPServer(port uint16) (tcps *TCPServer, err error) {
	tcpServ, err := net.Listen("tcp", ":"+strconv.Itoa(int(port)))
	if err != nil {
		return
	}
	tcps = &TCPServer{
		Listener:    tcpServ,
		addressbook: make(map[NetAddress]struct{}),
		// default handlers
		handlerMap: map[byte]func(net.Conn, []byte) error{
			'H': sendHostname,
			'P': tcps.sharePeers,
		},
	}

	// spawn listener
	go tcps.listen()
	return
}

// listen runs in the background, accepting incoming connections and serving
// them. listen will return after TCPServer.Close() is called, because the
// Accept() call will fail.
func (tcps *TCPServer) listen() {
	for {
		conn, err := tcps.Accept()
		if err != nil {
			return
		}
		// it is the handler's responsibility to close the connection
		go tcps.handleConn(conn)
	}
}

// handleConn reads header data from a connection, unmarshals the data
// structures it contains, and routes the data to other functions for
// processing.
// TODO: set deadlines?
func (tcps *TCPServer) handleConn(conn net.Conn) {
	defer conn.Close()
	var (
		msgHead []byte = make([]byte, 5)
		msgData []byte // length determined by msgHead
	)
	if n, err := conn.Read(msgHead); err != nil || n != 5 {
		// TODO: log error
		return
	}
	msgLen := DecUint64(msgHead[1:])
	if msgLen > maxMsgLen {
		// TODO: log error
		return
	}
	if msgLen > 0 {
		msgData = make([]byte, msgLen)
		if n, err := conn.Read(msgData); err != nil || uint64(n) != msgLen {
			// TODO: log error
			return
		}
	}

	// call registered handler for this message type
	if fn, ok := tcps.handlerMap[msgHead[0]]; ok {
		fn(conn, msgData)
		// TODO: log error
	}
	return
}

// sendHostname replies to the send with the sender's external IP.
func sendHostname(conn net.Conn, _ []byte) error {
	_, err := conn.Write([]byte(conn.RemoteAddr().String()))
	return err
}

// sharePeers transmits at most 'num' peers over the connection.
// TODO: choose random peers?
func (tcps *TCPServer) sharePeers(conn net.Conn, msgData []byte) error {
	if len(msgData) != 1 {
		return errors.New("invalid number of peers")
	}
	num := msgData[0]
	var addrs []NetAddress
	for addr := range tcps.addressbook {
		if num == 0 {
			break
		}
		addrs = append(addrs, addr)
		num--
	}
	_, err := conn.Write(Marshal(addrs))
	return err
}

// Ping returns whether a NetAddress is reachable. It accomplishes this by
// initiating a TCP connection and immediately closes it. This is pretty
// unsophisticated. I'll add a Pong later.
func (tcps *TCPServer) Ping(addr NetAddress) bool {
	conn, err := net.DialTimeout("tcp", addr.String(), timeout)
	if err != nil {
		return false
	}
	conn.Close()
	return true
}

// learnHostname learns the external IP of the TCPServer.
func (tcps *TCPServer) learnHostname(conn net.Conn) (err error) {
	// send hostname request
	if _, err = conn.Write([]byte{'H', 0, 0, 0, 0}); err != nil {
		return
	}
	// read response
	buf := make([]byte, 128)
	n, err := conn.Read(buf)
	if err != nil {
		return
	}
	// TODO: try to ping ourselves?
	host, portStr, err := net.SplitHostPort(string(buf[:n]))
	if err != nil {
		return
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return
	}
	tcps.myAddr = NetAddress{host, uint16(port)}
	return
}

// requestPeers queries a peer for additional peers, and adds any new peers to
// the address book.
func (tcps *TCPServer) requestPeers(conn net.Conn) (err error) {
	// request 10 peers
	if _, err = conn.Write([]byte{'P', 1, 0, 0, 0, 10}); err != nil {
		return
	}
	// read response
	buf := make([]byte, 1024)
	n, err := conn.Read(buf)
	if err != nil {
		return
	}
	var addrs []NetAddress
	if err = Unmarshal(buf[:n], &addrs); err != nil {
		return
	}
	// add peers
	// TODO: make sure we don't add ourself
	for _, addr := range addrs {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}
	return
}

// Bootstrap discovers the external IP of the TCPServer, requests peers from
// the initial peer list, and announces itself to those peers.
func (tcps *TCPServer) Bootstrap() (err error) {
	// populate initial peer list
	for _, addr := range BootstrapPeers {
		if tcps.Ping(addr) {
			tcps.addressbook[addr] = struct{}{}
		}
	}

	// learn hostname
	for addr := range tcps.addressbook {
		if addr.Call(tcps.learnHostname) == nil {
			break
		}
	}
	// request peers
	// TODO: maybe iterate until we have enough new peers?
	for addr := range tcps.addressbook {
		addr.Call(tcps.requestPeers)
	}
	// TODO: announce ourselves to new peers
	return
}
<|endoftext|>


<commit_before>package fakes

import (
	"time"

	sigar "github.com/cloudfoundry/gosigar"
)

type FakeSigar struct {
	LoadAverage    sigar.LoadAverage
	LoadAverageErr error

	Mem    sigar.Mem
	MemErr error

	Swap    sigar.Swap
	SwapErr error

	FileSystemUsage     sigar.FileSystemUsage
	FileSystemUsageErr  error
	FileSystemUsagePath string

	CollectCpuStatsCpuCh  chan sigar.Cpu
	CollectCpuStatsStopCh chan struct{}
}

func NewFakeSigar() *FakeSigar {
	return &FakeSigar{
		CollectCpuStatsCpuCh:  make(chan sigar.Cpu, 1),
		CollectCpuStatsStopCh: make(chan struct{}),
	}
}

func (f *FakeSigar) CollectCpuStats(collectionInterval time.Duration) (<-chan sigar.Cpu, chan<- struct{}) {
	samplesCh := make(chan sigar.Cpu, 1)
	stopCh := make(chan struct{})

	go func() {
		for {
			select {
			case cpuStat := <-f.CollectCpuStatsCpuCh:
				select {
				case samplesCh <- cpuStat:
				default:
					// Include default to avoid channel blocking
				}

			case <-f.CollectCpuStatsStopCh:
				return
			}
		}
	}()

	return samplesCh, stopCh
}

func (f *FakeSigar) GetLoadAverage() (sigar.LoadAverage, error) {
	return f.LoadAverage, f.LoadAverageErr
}

func (f *FakeSigar) GetMem() (sigar.Mem, error) {
	return f.Mem, f.MemErr
}

func (f *FakeSigar) GetSwap() (sigar.Swap, error) {
	return f.Swap, f.SwapErr
}

func (f *FakeSigar) GetFileSystemUsage(path string) (sigar.FileSystemUsage, error) {
	f.FileSystemUsagePath = path
	return f.FileSystemUsage, f.FileSystemUsageErr
}
<commit_msg>Add GetMemIgnoringCGroups to fake sigar interface<commit_after>package fakes

import (
	"time"

	sigar "github.com/cloudfoundry/gosigar"
)

type FakeSigar struct {
	LoadAverage    sigar.LoadAverage
	LoadAverageErr error

	Mem                sigar.Mem
	MemIgnoringCGroups sigar.Mem
	MemErr             error

	Swap    sigar.Swap
	SwapErr error

	FileSystemUsage     sigar.FileSystemUsage
	FileSystemUsageErr  error
	FileSystemUsagePath string

	CollectCpuStatsCpuCh  chan sigar.Cpu
	CollectCpuStatsStopCh chan struct{}
}

func NewFakeSigar() *FakeSigar {
	return &FakeSigar{
		CollectCpuStatsCpuCh:  make(chan sigar.Cpu, 1),
		CollectCpuStatsStopCh: make(chan struct{}),
	}
}

func (f *FakeSigar) CollectCpuStats(collectionInterval time.Duration) (<-chan sigar.Cpu, chan<- struct{}) {
	samplesCh := make(chan sigar.Cpu, 1)
	stopCh := make(chan struct{})

	go func() {
		for {
			select {
			case cpuStat := <-f.CollectCpuStatsCpuCh:
				select {
				case samplesCh <- cpuStat:
				default:
					// Include default to avoid channel blocking
				}

			case <-f.CollectCpuStatsStopCh:
				return
			}
		}
	}()

	return samplesCh, stopCh
}

func (f *FakeSigar) GetLoadAverage() (sigar.LoadAverage, error) {
	return f.LoadAverage, f.LoadAverageErr
}

func (f *FakeSigar) GetMem() (sigar.Mem, error) {
	return f.Mem, f.MemErr
}

func (f *FakeSigar) GetMemIgnoringCGroups() (sigar.Mem, error) {
	return f.MemIgnoringCGroups, f.MemErr
}

func (f *FakeSigar) GetSwap() (sigar.Swap, error) {
	return f.Swap, f.SwapErr
}

func (f *FakeSigar) GetFileSystemUsage(path string) (sigar.FileSystemUsage, error) {
	f.FileSystemUsagePath = path
	return f.FileSystemUsage, f.FileSystemUsageErr
}
<|endoftext|>


<commit_before>package comm

import (
	"bytes"
	"errors"
	"fmt"
	"io"
)

const (
	BufferSize = 16
)

var (
	EOFError = errors.New("Buffer is exhausted")
)

type Stream struct {
	data   []byte
	Reader *bytes.Buffer
	conn   io.Reader
}

type errorTrampolineFunc func() (errorTrampolineFunc, error)

func NewStream(conn io.Reader) (*Stream, error) {
	data := make([]byte, BufferSize)

	n, err := conn.Read(data)
	if err != nil {
		return nil, err
	}

	fmt.Printf("Read %d bytes\n", n)

	stream := &Stream{
		data:   data,
		Reader: bytes.NewBuffer(data[0:n]),
		conn:   conn,
	}

	return stream, nil
}

func (s *Stream) ReadWith(fn func() error) error {
	return errorTrampoline(s.readWith(fn))
}

func (s *Stream) readWith(fn func() error) (errorTrampolineFunc, error) {
	bytesBefore := make([]byte, s.Reader.Len())
	copy(bytesBefore, s.Reader.Bytes())

	err := fn()
	if err == nil {
		return nil, nil
	}

	n, err := s.conn.Read(s.data)
	if err != nil {
		return nil, err
	}

	fmt.Printf("Read %d bytes\n", n)

	s.Reader = bytes.NewBuffer(
		append(bytesBefore, s.data[0:n]...),
	)

	return func() (errorTrampolineFunc, error) {
		return s.readWith(fn)
	}, nil
}

func errorTrampoline(fn errorTrampolineFunc, err error) error {
	if err != nil {
		return err
	}

	if fn == nil {
		return nil
	}

	for {
		fn, err = fn()
		if err != nil {
			return err
		}
		if fn == nil {
			return nil
		}
	}
}
<commit_msg>Restore normal buffer size<commit_after>package comm

import (
	"bytes"
	"errors"
	"fmt"
	"io"
)

const (
	BufferSize = 4096
)

var (
	EOFError = errors.New("Buffer is exhausted")
)

type Stream struct {
	data   []byte
	Reader *bytes.Buffer
	conn   io.Reader
}

type errorTrampolineFunc func() (errorTrampolineFunc, error)

func NewStream(conn io.Reader) (*Stream, error) {
	data := make([]byte, BufferSize)

	n, err := conn.Read(data)
	if err != nil {
		return nil, err
	}

	fmt.Printf("Read %d bytes\n", n)

	stream := &Stream{
		data:   data,
		Reader: bytes.NewBuffer(data[0:n]),
		conn:   conn,
	}

	return stream, nil
}

func (s *Stream) ReadWith(fn func() error) error {
	return errorTrampoline(s.readWith(fn))
}

func (s *Stream) readWith(fn func() error) (errorTrampolineFunc, error) {
	bytesBefore := make([]byte, s.Reader.Len())
	copy(bytesBefore, s.Reader.Bytes())

	err := fn()
	if err == nil {
		return nil, nil
	}

	n, err := s.conn.Read(s.data)
	if err != nil {
		return nil, err
	}

	fmt.Printf("Read %d bytes\n", n)

	s.Reader = bytes.NewBuffer(
		append(bytesBefore, s.data[0:n]...),
	)

	return func() (errorTrampolineFunc, error) {
		return s.readWith(fn)
	}, nil
}

func errorTrampoline(fn errorTrampolineFunc, err error) error {
	if err != nil {
		return err
	}

	if fn == nil {
		return nil
	}

	for {
		fn, err = fn()
		if err != nil {
			return err
		}
		if fn == nil {
			return nil
		}
	}
}
<|endoftext|>


<commit_before>package testutils

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

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega/gexec"
)

func BuildExecutable() error {
	return BuildExecutableForArch("")
}

func BuildExecutableForArch(arch string) error {
	buildArg := "./../bin/build"
	if arch != "" {
		buildArg = buildArg + "-" + arch
	}

	session, err := RunCommand(buildArg)
	if session.ExitCode() != 0 {
		return fmt.Errorf("Failed to build bosh:\nstdout:\n%s\nstderr:\n%s", session.Out.Contents(), session.Err.Contents())
	}

	return err
}

func RunCommand(cmd string, args ...string) (*gexec.Session, error) {
	return RunComplexCommand(exec.Command(cmd, args...))
}

func RunComplexCommand(cmd *exec.Cmd) (*gexec.Session, error) {
	session, err := gexec.Start(cmd, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	if err != nil {
		return nil, err
	}

	session.Wait(60 * time.Second)

	return session, nil
}
<commit_msg>Bump binary build timeout for slow CI workers<commit_after>package testutils

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

	"github.com/onsi/ginkgo"
	"github.com/onsi/gomega/gexec"
)

func BuildExecutable() error {
	return BuildExecutableForArch("")
}

func BuildExecutableForArch(arch string) error {
	buildArg := "./../bin/build"
	if arch != "" {
		buildArg = buildArg + "-" + arch
	}

	session, err := RunCommand(buildArg)
	if session.ExitCode() != 0 {
		return fmt.Errorf("Failed to build bosh:\nstdout:\n%s\nstderr:\n%s", session.Out.Contents(), session.Err.Contents())
	}

	return err
}

func RunCommand(cmd string, args ...string) (*gexec.Session, error) {
	return RunComplexCommand(exec.Command(cmd, args...))
}

func RunComplexCommand(cmd *exec.Cmd) (*gexec.Session, error) {
	session, err := gexec.Start(cmd, ginkgo.GinkgoWriter, ginkgo.GinkgoWriter)
	if err != nil {
		return nil, err
	}

	session.Wait(120 * time.Second)

	return session, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2021 The Vitess Authors.

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

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

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

package command

import (
	"time"

	"github.com/spf13/cobra"

	"vitess.io/vitess/go/cmd/vtctldclient/cli"
	"vitess.io/vitess/go/protoutil"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/topo/topoproto"

	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
)

var (
	// InitShardPrimary makes an InitShardPrimary gRPC call to a vtctld.
	InitShardPrimary = &cobra.Command{
		Use:  "InitShardPrimary",
		Args: cobra.ExactArgs(2),
		RunE: commandInitShardPrimary,
	}
)

var initShardPrimaryOptions = struct {
	WaitReplicasTimeout time.Duration
	Force               bool
}{}

func commandInitShardPrimary(cmd *cobra.Command, args []string) error {
	keyspace, shard, err := topoproto.ParseKeyspaceShard(cmd.Flags().Arg(0))
	if err != nil {
		return err
	}

	tabletAlias, err := topoproto.ParseTabletAlias(cmd.Flags().Arg(1))
	if err != nil {
		return err
	}

	cli.FinishedParsing(cmd)

	resp, err := client.InitShardPrimary(commandCtx, &vtctldatapb.InitShardPrimaryRequest{
		Keyspace:                keyspace,
		Shard:                   shard,
		PrimaryElectTabletAlias: tabletAlias,
		WaitReplicasTimeout:     protoutil.DurationToProto(initShardPrimaryOptions.WaitReplicasTimeout),
		Force:                   initShardPrimaryOptions.Force,
	})

	for _, event := range resp.Events {
		log.Infof("%v", event)
	}

	return err
}

func init() {
	InitShardPrimary.Flags().DurationVar(&initShardPrimaryOptions.WaitReplicasTimeout, "wait-replicas-timeout", 30*time.Second, "time to wait for replicas to catch up in reparenting")
	InitShardPrimary.Flags().BoolVar(&initShardPrimaryOptions.Force, "force", false, "will force the reparent even if the provided tablet is not a master or the shard master")
	Root.AddCommand(InitShardPrimary)
}
<commit_msg>Add CLI entrypoint for `TabletExternallyReparented`<commit_after>/*
Copyright 2021 The Vitess Authors.

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

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

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

package command

import (
	"fmt"
	"time"

	"github.com/spf13/cobra"

	"vitess.io/vitess/go/cmd/vtctldclient/cli"
	"vitess.io/vitess/go/protoutil"
	"vitess.io/vitess/go/vt/log"
	"vitess.io/vitess/go/vt/topo/topoproto"

	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
)

var (
	// InitShardPrimary makes an InitShardPrimary gRPC call to a vtctld.
	InitShardPrimary = &cobra.Command{
		Use:  "InitShardPrimary",
		Args: cobra.ExactArgs(2),
		RunE: commandInitShardPrimary,
	}
	// TabletExternallyReparented makes a TabletExternallyReparented gRPC call
	// to a vtctld.
	TabletExternallyReparented = &cobra.Command{
		Use:  "TabletExternallyReparented ALIAS",
		Args: cobra.ExactArgs(1),
		RunE: commandTabletExternallyReparented,
	}
)

var initShardPrimaryOptions = struct {
	WaitReplicasTimeout time.Duration
	Force               bool
}{}

func commandInitShardPrimary(cmd *cobra.Command, args []string) error {
	keyspace, shard, err := topoproto.ParseKeyspaceShard(cmd.Flags().Arg(0))
	if err != nil {
		return err
	}

	tabletAlias, err := topoproto.ParseTabletAlias(cmd.Flags().Arg(1))
	if err != nil {
		return err
	}

	cli.FinishedParsing(cmd)

	resp, err := client.InitShardPrimary(commandCtx, &vtctldatapb.InitShardPrimaryRequest{
		Keyspace:                keyspace,
		Shard:                   shard,
		PrimaryElectTabletAlias: tabletAlias,
		WaitReplicasTimeout:     protoutil.DurationToProto(initShardPrimaryOptions.WaitReplicasTimeout),
		Force:                   initShardPrimaryOptions.Force,
	})

	for _, event := range resp.Events {
		log.Infof("%v", event)
	}

	return err
}

func commandTabletExternallyReparented(cmd *cobra.Command, args []string) error {
	alias, err := topoproto.ParseTabletAlias(cmd.Flags().Arg(0))
	if err != nil {
		return err
	}

	resp, err := client.TabletExternallyReparented(commandCtx, &vtctldatapb.TabletExternallyReparentedRequest{
		Tablet: alias,
	})
	if err != nil {
		return err
	}

	data, err := cli.MarshalJSON(resp)
	if err != nil {
		return err
	}

	fmt.Printf("%s\n", data)

	return nil
}

func init() {
	InitShardPrimary.Flags().DurationVar(&initShardPrimaryOptions.WaitReplicasTimeout, "wait-replicas-timeout", 30*time.Second, "time to wait for replicas to catch up in reparenting")
	InitShardPrimary.Flags().BoolVar(&initShardPrimaryOptions.Force, "force", false, "will force the reparent even if the provided tablet is not a master or the shard master")
	Root.AddCommand(InitShardPrimary)

	Root.AddCommand(TabletExternallyReparented)
}
<|endoftext|>


<commit_before>/*
Copyright 2022 The Vitess Authors.

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

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

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

package misc

import (
	"fmt"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/test/endtoend/cluster"
	"vitess.io/vitess/go/test/endtoend/utils"
)

func start(t *testing.T) (utils.MySQLCompare, func()) {
	mcmp, err := utils.NewMySQLCompare(t, vtParams, mysqlParams)
	require.NoError(t, err)

	deleteAll := func() {
		tables := []string{"t1"}
		for _, table := range tables {
			_, _ = mcmp.ExecAndIgnore("delete from " + table)
		}
	}

	deleteAll()

	return mcmp, func() {
		deleteAll()
		mcmp.Close()
		cluster.PanicHandler(t)
	}
}

func TestBitVals(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	mcmp.Exec("insert into t1(id1, id2) values (0,0)")

	mcmp.AssertMatches(`select b'1001', 0x9, B'010011011010'`, `[[VARBINARY("\t") VARBINARY("\t") VARBINARY("\x04\xda")]]`)
	mcmp.AssertMatches(`select b'1001', 0x9, B'010011011010' from t1`, `[[VARBINARY("\t") VARBINARY("\t") VARBINARY("\x04\xda")]]`)
	mcmp.AssertMatchesNoCompare(`select 1 + b'1001', 2 + 0x9, 3 + B'010011011010'`, `[[INT64(10) UINT64(11) INT64(1245)]]`, `[[UINT64(10) UINT64(11) UINT64(1245)]]`)
	mcmp.AssertMatchesNoCompare(`select 1 + b'1001', 2 + 0x9, 3 + B'010011011010' from t1`, `[[INT64(10) UINT64(11) INT64(1245)]]`, `[[UINT64(10) UINT64(11) UINT64(1245)]]`)
}

func TestHexVals(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	mcmp.Exec("insert into t1(id1, id2) values (0,0)")

	mcmp.AssertMatches(`select x'09', 0x9`, `[[VARBINARY("\t") VARBINARY("\t")]]`)
	mcmp.AssertMatches(`select X'09', 0x9 from t1`, `[[VARBINARY("\t") VARBINARY("\t")]]`)
	mcmp.AssertMatches(`select 1 + x'09', 2 + 0x9`, `[[UINT64(10) UINT64(11)]]`)
	mcmp.AssertMatches(`select 1 + X'09', 2 + 0x9 from t1`, `[[UINT64(10) UINT64(11)]]`)
}

func TestDateTimeTimestampVals(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	mcmp.AssertMatches(`select date'2022-08-03'`, `[[DATE("2022-08-03")]]`)
	mcmp.AssertMatches(`select time'12:34:56'`, `[[TIME("12:34:56")]]`)
	mcmp.AssertMatches(`select timestamp'2012-12-31 11:30:45'`, `[[DATETIME("2012-12-31 11:30:45")]]`)
}

func TestInvalidDateTimeTimestampVals(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	_, err := mcmp.ExecAllowAndCompareError(`select date'2022'`)
	require.Error(t, err)
	_, err = mcmp.ExecAllowAndCompareError(`select time'12:34:56:78'`)
	require.Error(t, err)
	_, err = mcmp.ExecAllowAndCompareError(`select timestamp'2022'`)
	require.Error(t, err)
}

func TestQueryTimeoutWithDual(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	_, err := utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 */ sleep(0.04) from dual")
	assert.NoError(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 */ sleep(0.24) from dual")
	assert.Error(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "set @@session.query_timeout=20")
	require.NoError(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 */ sleep(0.04) from dual")
	assert.Error(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 */ sleep(0.01) from dual")
	assert.NoError(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=500 */ sleep(0.24) from dual")
	assert.NoError(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=10 */ sleep(0.04) from dual")
	assert.Error(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=10 */ sleep(0.001) from dual")
	assert.NoError(t, err)
}

func TestQueryTimeoutWithTables(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	// unsharded
	utils.Exec(t, mcmp.VtConn, "insert /*vt+ QUERY_TIMEOUT_MS=1000 */ into uks.unsharded(id1) values (1),(2),(3),(4),(5)")
	for i := 0; i < 12; i++ {
		utils.Exec(t, mcmp.VtConn, "insert /*vt+ QUERY_TIMEOUT_MS=1000 */ into uks.unsharded(id1) select id1+5 from uks.unsharded")
	}

	utils.Exec(t, mcmp.VtConn, "select count(*) from uks.unsharded where id1 > 31")
	utils.Exec(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=100 */ count(*) from uks.unsharded where id1 > 31")

	// the query usually takes more than 5ms to return. So this should fail.
	_, err := utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=1 */ count(*) from uks.unsharded where id1 > 31")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "context deadline exceeded (errno 1317) (sqlstate 70100)")

	// sharded
	for i := 0; i < 300000; i += 1000 {
		var str strings.Builder
		for j := 1; j <= 1000; j++ {
			if j == 1 {
				str.WriteString(fmt.Sprintf("(%d)", i*1000+j))
				continue
			}
			str.WriteString(fmt.Sprintf(",(%d)", i*1000+j))
		}
		utils.Exec(t, mcmp.VtConn, fmt.Sprintf("insert /*vt+ QUERY_TIMEOUT_MS=1000 */ into t1(id1) values %s", str.String()))
	}

	utils.Exec(t, mcmp.VtConn, "select count(*) from t1 where id1 > 31")
	utils.Exec(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=100 */ count(*) from t1 where id1 > 31")

	// the query usually takes more than 5ms to return. So this should fail.
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=1 */ count(*) from t1 where id1 > 31")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "context deadline exceeded (errno 1317) (sqlstate 70100)")
}
<commit_msg>test: deflake TestQueryTimeoutWithTables (#11772)<commit_after>/*
Copyright 2022 The Vitess Authors.

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

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

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

package misc

import (
	"fmt"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"vitess.io/vitess/go/test/endtoend/cluster"
	"vitess.io/vitess/go/test/endtoend/utils"
)

func start(t *testing.T) (utils.MySQLCompare, func()) {
	mcmp, err := utils.NewMySQLCompare(t, vtParams, mysqlParams)
	require.NoError(t, err)

	deleteAll := func() {
		tables := []string{"t1"}
		for _, table := range tables {
			_, _ = mcmp.ExecAndIgnore("delete from " + table)
		}
	}

	deleteAll()

	return mcmp, func() {
		deleteAll()
		mcmp.Close()
		cluster.PanicHandler(t)
	}
}

func TestBitVals(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	mcmp.Exec("insert into t1(id1, id2) values (0,0)")

	mcmp.AssertMatches(`select b'1001', 0x9, B'010011011010'`, `[[VARBINARY("\t") VARBINARY("\t") VARBINARY("\x04\xda")]]`)
	mcmp.AssertMatches(`select b'1001', 0x9, B'010011011010' from t1`, `[[VARBINARY("\t") VARBINARY("\t") VARBINARY("\x04\xda")]]`)
	mcmp.AssertMatchesNoCompare(`select 1 + b'1001', 2 + 0x9, 3 + B'010011011010'`, `[[INT64(10) UINT64(11) INT64(1245)]]`, `[[UINT64(10) UINT64(11) UINT64(1245)]]`)
	mcmp.AssertMatchesNoCompare(`select 1 + b'1001', 2 + 0x9, 3 + B'010011011010' from t1`, `[[INT64(10) UINT64(11) INT64(1245)]]`, `[[UINT64(10) UINT64(11) UINT64(1245)]]`)
}

func TestHexVals(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	mcmp.Exec("insert into t1(id1, id2) values (0,0)")

	mcmp.AssertMatches(`select x'09', 0x9`, `[[VARBINARY("\t") VARBINARY("\t")]]`)
	mcmp.AssertMatches(`select X'09', 0x9 from t1`, `[[VARBINARY("\t") VARBINARY("\t")]]`)
	mcmp.AssertMatches(`select 1 + x'09', 2 + 0x9`, `[[UINT64(10) UINT64(11)]]`)
	mcmp.AssertMatches(`select 1 + X'09', 2 + 0x9 from t1`, `[[UINT64(10) UINT64(11)]]`)
}

func TestDateTimeTimestampVals(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	mcmp.AssertMatches(`select date'2022-08-03'`, `[[DATE("2022-08-03")]]`)
	mcmp.AssertMatches(`select time'12:34:56'`, `[[TIME("12:34:56")]]`)
	mcmp.AssertMatches(`select timestamp'2012-12-31 11:30:45'`, `[[DATETIME("2012-12-31 11:30:45")]]`)
}

func TestInvalidDateTimeTimestampVals(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	_, err := mcmp.ExecAllowAndCompareError(`select date'2022'`)
	require.Error(t, err)
	_, err = mcmp.ExecAllowAndCompareError(`select time'12:34:56:78'`)
	require.Error(t, err)
	_, err = mcmp.ExecAllowAndCompareError(`select timestamp'2022'`)
	require.Error(t, err)
}

func TestQueryTimeoutWithDual(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	_, err := utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 */ sleep(0.04) from dual")
	assert.NoError(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 */ sleep(0.24) from dual")
	assert.Error(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "set @@session.query_timeout=20")
	require.NoError(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 */ sleep(0.04) from dual")
	assert.Error(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 */ sleep(0.01) from dual")
	assert.NoError(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=500 */ sleep(0.24) from dual")
	assert.NoError(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=10 */ sleep(0.04) from dual")
	assert.Error(t, err)
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=10 */ sleep(0.001) from dual")
	assert.NoError(t, err)
}

func TestQueryTimeoutWithTables(t *testing.T) {
	mcmp, closer := start(t)
	defer closer()

	// unsharded
	utils.Exec(t, mcmp.VtConn, "insert /*vt+ QUERY_TIMEOUT_MS=1000 */ into uks.unsharded(id1) values (1),(2),(3),(4),(5)")
	for i := 0; i < 12; i++ {
		utils.Exec(t, mcmp.VtConn, "insert /*vt+ QUERY_TIMEOUT_MS=1000 */ into uks.unsharded(id1) select id1+5 from uks.unsharded")
	}

	utils.Exec(t, mcmp.VtConn, "select count(*) from uks.unsharded where id1 > 31")
	utils.Exec(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=100 */ count(*) from uks.unsharded where id1 > 31")

	// the query usually takes more than 5ms to return. So this should fail.
	_, err := utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=1 */ count(*) from uks.unsharded where id1 > 31")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "context deadline exceeded")
	assert.Contains(t, err.Error(), "(errno 1317) (sqlstate 70100)")

	// sharded
	for i := 0; i < 300000; i += 1000 {
		var str strings.Builder
		for j := 1; j <= 1000; j++ {
			if j == 1 {
				str.WriteString(fmt.Sprintf("(%d)", i*1000+j))
				continue
			}
			str.WriteString(fmt.Sprintf(",(%d)", i*1000+j))
		}
		utils.Exec(t, mcmp.VtConn, fmt.Sprintf("insert /*vt+ QUERY_TIMEOUT_MS=1000 */ into t1(id1) values %s", str.String()))
	}

	utils.Exec(t, mcmp.VtConn, "select count(*) from t1 where id1 > 31")
	utils.Exec(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=100 */ count(*) from t1 where id1 > 31")

	// the query usually takes more than 5ms to return. So this should fail.
	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ PLANNER=gen4 QUERY_TIMEOUT_MS=1 */ count(*) from t1 where id1 > 31")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "context deadline exceeded")
	assert.Contains(t, err.Error(), "(errno 1317) (sqlstate 70100)")
}
<|endoftext|>


<commit_before>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 log

// BlackHole logs and does nothing. An empty struct. IsDebug() and IsInfo() will
// always return true.
type BlackHole struct{}

// New returns a new Logger that has this logger's context plus the given context
func (l BlackHole) New(ctx ...interface{}) Logger { return BlackHole{} }

// Debug logs a debug entry.
func (l BlackHole) Debug(msg string, args ...interface{}) {}

// Info logs an info entry.
func (l BlackHole) Info(msg string, args ...interface{}) {}

// Fatal logs a fatal entry then panics.
func (l BlackHole) Fatal(msg string, args ...interface{}) { panic("exit due to fatal error: " + msg) }

// IsDebug determines if this logger logs a debug statement. Returns always true.
func (l BlackHole) IsDebug() bool { return true }

// IsInfo determines if this logger logs an info statement. Returns always true.
func (l BlackHole) IsInfo() bool { return true }

// SetLevel sets the level of this logger.
func (l BlackHole) SetLevel(level int) {}
<commit_msg>utils/log: BlackHole can now enable debug and/or info logging<commit_after>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// 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 log

// BlackHole logs and does nothing. An empty struct. IsDebug() and IsInfo() will
// always return true.
type BlackHole struct {
	EnableDebug bool
	EnableInfo  bool
}

// NewBlackHole creates a new black hole logger where debug and info logging
// is enabled.
func NewBlackHole() BlackHole {
	return BlackHole{true, true}
}

// New returns a new Logger that has this logger's context plus the given context
func (l BlackHole) New(ctx ...interface{}) Logger { return BlackHole{} }

// Debug logs a debug entry. Noop.
func (l BlackHole) Debug(msg string, args ...interface{}) {}

// Info logs an info entry. Noop.
func (l BlackHole) Info(msg string, args ...interface{}) {}

// Fatal logs a fatal entry then panics.
func (l BlackHole) Fatal(msg string, args ...interface{}) { panic("exit due to fatal error: " + msg) }

// IsDebug determines if this logger logs a debug statement. Returns always true.
func (l BlackHole) IsDebug() bool { return l.EnableDebug }

// IsInfo determines if this logger logs an info statement. Returns always true.
func (l BlackHole) IsInfo() bool { return l.EnableInfo }

// SetLevel sets the level of this logger. Noop.
func (l BlackHole) SetLevel(level int) {}
<|endoftext|>


<commit_before>package timing

import "time"

// Controls whether library calls should be mocked, or whether we should use the standard Go time library.
// If we're in Mock Mode, then time does not pass as normal, but only progresses when Elapse is called.
// False by default, indicating that we just call through to standard Go functions.
var MockMode bool = false
var currentTimeMock time.Time = time.Unix(0, 0)
var mockTimers []*mockTimer = make([]*mockTimer, 0)

// Interface over Golang's built-in Timers, allowing them to be swapped out for mocked timers.
type Timer interface {
	// Returns a channel which sends the current time immediately when the timer expires.
	// Equivalent to time.Timer.C; however, we have to use a method here instead of a member since this is an interface.
	C() <-chan time.Time

	// Resets the timer such that it will expire in duration 'd' after the current time.
	// Returns true if the timer had been active, and false if it had expired or been stopped.
	Reset(d time.Duration) bool

	// Stops the timer, preventing it from firing.
	// Returns true if the timer had been active, and false if it had expired or been stopped.
	Stop() bool
}

// Implementation of Timer that just wraps time.Timer.
type realTimer struct {
	*time.Timer
}

func (t *realTimer) C() <-chan time.Time {
	return t.Timer.C
}

func (t *realTimer) Reset(d time.Duration) bool {
	return t.Timer.Reset(d)
}

func (t *realTimer) Stop() bool {
	return t.Timer.Stop()
}

// Implementation of Timer that mocks time.Timer, firing when the total elapsed time (as controlled by Elapse)
// exceeds the duration specified when the timer was constructed.
type mockTimer struct {
	EndTime time.Time
	Chan    chan time.Time
}

func (t *mockTimer) C() <-chan time.Time {
	return t.Chan
}

func (t *mockTimer) Reset(d time.Duration) bool {
	wasActive := removeMockTimer(t)

	t.EndTime = currentTimeMock.Add(d)
	if d > 0 {
		mockTimers = append(mockTimers, t)
	} else {
		// The new timer has an expiry time of 0.
		// Fire it right away, and don't bother tracking it.
		t.Chan <- currentTimeMock
	}

	return wasActive
}

func (t *mockTimer) Stop() bool {
	return removeMockTimer(t)
}

// Creates a new Timer; either a wrapper around a standard Go time.Timer, or a mocked-out Timer,
// depending on whether MockMode is set.
func NewTimer(d time.Duration) Timer {
	if MockMode {
		t := mockTimer{currentTimeMock.Add(d), make(chan time.Time)}
		if d == 0 {
			t.Chan <- currentTimeMock
		} else {
			mockTimers = append(mockTimers, &t)
		}
		return &t
	} else {
		return &realTimer{time.NewTimer(d)}
	}
}

// See built-in time.After() function.
func After(d time.Duration) <-chan time.Time {
	return NewTimer(d).C()
}

// See built-in time.AfterFunc() function.
func AfterFunc(d time.Duration, f func()) Timer {
	t := NewTimer(d)
	go func() {
		if _, ok := <-t.C(); ok {
			f()
		}
	}()
	return t
}

// See built-in time.Sleep() function.
func Sleep(d time.Duration) {
	<-After(d)
}

// Increment the current time by the given Duration.
// This function can only be called in Mock Mode, otherwise we will panic.
func Elapse(d time.Duration) {
	requireMockMode()
	fired := make([]*mockTimer, 0)
	currentTimeMock = currentTimeMock.Add(d)

	// Fire any timers whose time has come up.
	for _, t := range mockTimers {
		if currentTimeMock.After(t.EndTime) {
			t.Chan <- currentTimeMock
			fired = append(fired, t)
		}
	}

	// Stop tracking any fired timers.
	remainingTimers := make([]*mockTimer, 0)
loop:
	for _, t := range mockTimers {
		for _, fired_t := range fired {
			if t == fired_t {
				continue loop
			}
		}

		remainingTimers = append(remainingTimers, t)
	}

	mockTimers = remainingTimers
}

// Returns the current time.
// If Mock Mode is set, this will be the sum of all Durations passed into Elapse calls;
// otherwise it will be the true system time.
func Now() time.Time {
	if MockMode {
		return currentTimeMock
	} else {
		return time.Now()
	}
}

// Shortcut method to enforce that Mock Mode is enabled.
func requireMockMode() {
	if !MockMode {
		panic("This method requires MockMode to be enabled")
	}
}

// Utility method to remove a mockTimer from the list of outstanding timers.
func removeMockTimer(t *mockTimer) bool {
	// First, find the index of the timer in our list.
	found := false
	var idx int
	var elt *mockTimer
	for idx, elt = range mockTimers {
		if elt == t {
			found = true
			break
		}
	}

	if found {
		// We found the given timer. Remove it.
		if idx == len(mockTimers)-1 {
			mockTimers = mockTimers[:idx]
		} else {
			mockTimers = append(mockTimers[:idx], mockTimers[idx+1])
		}
		return true
	} else {
		// The timer was not present, indicating that it was already expired.
		return false
	}
}
<commit_msg>Fix off-by-one error with mock timers.  Prevent time.Elapse from blocking.<commit_after>package timing

import "time"

// Controls whether library calls should be mocked, or whether we should use the standard Go time library.
// If we're in Mock Mode, then time does not pass as normal, but only progresses when Elapse is called.
// False by default, indicating that we just call through to standard Go functions.
var MockMode bool = false
var currentTimeMock time.Time = time.Unix(0, 0)
var mockTimers []*mockTimer = make([]*mockTimer, 0)

// Interface over Golang's built-in Timers, allowing them to be swapped out for mocked timers.
type Timer interface {
	// Returns a channel which sends the current time immediately when the timer expires.
	// Equivalent to time.Timer.C; however, we have to use a method here instead of a member since this is an interface.
	C() <-chan time.Time

	// Resets the timer such that it will expire in duration 'd' after the current time.
	// Returns true if the timer had been active, and false if it had expired or been stopped.
	Reset(d time.Duration) bool

	// Stops the timer, preventing it from firing.
	// Returns true if the timer had been active, and false if it had expired or been stopped.
	Stop() bool
}

// Implementation of Timer that just wraps time.Timer.
type realTimer struct {
	*time.Timer
}

func (t *realTimer) C() <-chan time.Time {
	return t.Timer.C
}

func (t *realTimer) Reset(d time.Duration) bool {
	return t.Timer.Reset(d)
}

func (t *realTimer) Stop() bool {
	return t.Timer.Stop()
}

// Implementation of Timer that mocks time.Timer, firing when the total elapsed time (as controlled by Elapse)
// exceeds the duration specified when the timer was constructed.
type mockTimer struct {
	EndTime time.Time
	Chan    chan time.Time
}

func (t *mockTimer) C() <-chan time.Time {
	return t.Chan
}

func (t *mockTimer) Reset(d time.Duration) bool {
	wasActive := removeMockTimer(t)

	t.EndTime = currentTimeMock.Add(d)
	if d > 0 {
		mockTimers = append(mockTimers, t)
	} else {
		// The new timer has an expiry time of 0.
		// Fire it right away, and don't bother tracking it.
		t.Chan <- currentTimeMock
	}

	return wasActive
}

func (t *mockTimer) Stop() bool {
	return removeMockTimer(t)
}

// Creates a new Timer; either a wrapper around a standard Go time.Timer, or a mocked-out Timer,
// depending on whether MockMode is set.
func NewTimer(d time.Duration) Timer {
	if MockMode {
		t := mockTimer{currentTimeMock.Add(d), make(chan time.Time, 1)}
		if d == 0 {
			t.Chan <- currentTimeMock
		} else {
			mockTimers = append(mockTimers, &t)
		}
		return &t
	} else {
		return &realTimer{time.NewTimer(d)}
	}
}

// See built-in time.After() function.
func After(d time.Duration) <-chan time.Time {
	return NewTimer(d).C()
}

// See built-in time.AfterFunc() function.
func AfterFunc(d time.Duration, f func()) Timer {
	t := NewTimer(d)
	go func() {
		if _, ok := <-t.C(); ok {
			f()
		}
	}()
	return t
}

// See built-in time.Sleep() function.
func Sleep(d time.Duration) {
	<-After(d)
}

// Increment the current time by the given Duration.
// This function can only be called in Mock Mode, otherwise we will panic.
func Elapse(d time.Duration) {
	requireMockMode()
	fired := make([]*mockTimer, 0)
	currentTimeMock = currentTimeMock.Add(d)

	// Fire any timers whose time has come up.
	for _, t := range mockTimers {
		if !t.EndTime.Before(currentTimeMock) {
			t.Chan <- currentTimeMock
			fired = append(fired, t)
		}
	}

	// Stop tracking any fired timers.
	remainingTimers := make([]*mockTimer, 0)
loop:
	for _, t := range mockTimers {
		for _, fired_t := range fired {
			if t == fired_t {
				continue loop
			}
		}

		remainingTimers = append(remainingTimers, t)
	}

	mockTimers = remainingTimers
}

// Returns the current time.
// If Mock Mode is set, this will be the sum of all Durations passed into Elapse calls;
// otherwise it will be the true system time.
func Now() time.Time {
	if MockMode {
		return currentTimeMock
	} else {
		return time.Now()
	}
}

// Shortcut method to enforce that Mock Mode is enabled.
func requireMockMode() {
	if !MockMode {
		panic("This method requires MockMode to be enabled")
	}
}

// Utility method to remove a mockTimer from the list of outstanding timers.
func removeMockTimer(t *mockTimer) bool {
	// First, find the index of the timer in our list.
	found := false
	var idx int
	var elt *mockTimer
	for idx, elt = range mockTimers {
		if elt == t {
			found = true
			break
		}
	}

	if found {
		// We found the given timer. Remove it.
		if idx == len(mockTimers)-1 {
			mockTimers = mockTimers[:idx]
		} else {
			mockTimers = append(mockTimers[:idx], mockTimers[idx+1])
		}
		return true
	} else {
		// The timer was not present, indicating that it was already expired.
		return false
	}
}
<|endoftext|>


<commit_before>package prose

import (
	"encoding/json"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/assert"
)

var testdata = "testdata"

func makeDoc(text string) (*Document, error) {
	return NewDocument(
		text,
		WithSegmentation(false),
		WithTagging(false),
		WithExtraction(false))
}

func getTokenText(doc *Document) []string {
	observed := []string{}
	for _, tok := range doc.Tokens() {
		observed = append(observed, tok.Text)
	}
	return observed
}

func getWordData(file string) ([]string, [][]string) {
	in := readDataFile(filepath.Join(testdata, "treebank_sents.json"))
	out := readDataFile(filepath.Join(testdata, file))

	input := []string{}
	output := [][]string{}

	checkError(json.Unmarshal(in, &input))
	checkError(json.Unmarshal(out, &output))

	return input, output
}

func getWordBenchData() []string {
	in := readDataFile(filepath.Join(testdata, "treebank_sents.json"))
	input := []string{}
	checkError(json.Unmarshal(in, &input))
	return input
}

func TestTokenizationSimple(t *testing.T) {
	doc, _ := makeDoc("Vale is a natural language linter that supports plain text, markup (Markdown, reStructuredText, AsciiDoc, and HTML), and source code comments. Vale doesn't attempt to offer a one-size-fits-all collection of rules—instead, it strives to make customization as easy as possible.")
	expected := []string{
		"Vale", "is", "a", "natural", "language", "linter", "that", "supports",
		"plain", "text", ",", "markup", "(", "Markdown", ",", "reStructuredText",
		",", "AsciiDoc", ",", "and", "HTML", ")", ",", "and", "source",
		"code", "comments", ".", "Vale", "does", "n't", "attempt", "to",
		"offer", "a", "one-size-fits-all", "collection", "of", "rules—instead",
		",", "it", "strives", "to", "make", "customization", "as", "easy", "as",
		"possible", "."}
	assert.Equal(t, expected, getTokenText(doc))
}

func TestTokenizationTreebank(t *testing.T) {
	input, output := getWordData("treebank_words.json")
	for i, s := range input {
		doc, _ := makeDoc(s)
		assert.Equal(t, output[i], getTokenText(doc))
	}
}

func TestTokenizationWeb(t *testing.T) {
	web := `Independent of current body composition, IGF-I levels at 5 yr were significantly
            associated with rate of weight gain between 0-2 yr (beta=0.19; P<0.0005);
            and children who showed postnatal catch-up growth (i.e. those who showed gains in
            weight or length between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other
			children (P=0.02; http://univ.edu.es/study.html) [20-22].`
	expected := []string{"Independent", "of", "current", "body", "composition", ",", "IGF-I",
		"levels", "at", "5", "yr", "were", "significantly", "associated", "with", "rate", "of",
		"weight", "gain", "between", "0-2", "yr", "(", "beta=0.19", ";", "P<0.0005", ")", ";",
		"and", "children", "who", "showed", "postnatal", "catch-up", "growth", "(", "i.e.", "those",
		"who", "showed", "gains", "in", "weight", "or", "length", "between", "0-2", "yr", "by",
		">0.67", "SD", "score", ")", "had", "higher", "IGF-I", "levels", "than", "other", "children",
		"(", "P=0.02", ";", "http://univ.edu.es/study.html", ")", "[", "20-22", "]", "."}
	doc, _ := makeDoc(web)
	assert.Equal(t, expected, getTokenText(doc))
}

func TestTokenizationWebParagraph(t *testing.T) {
	web := `Independent of current body composition, IGF-I levels at 5 yr were significantly
            associated with rate of weight gain between 0-2 yr (beta=0.19; P<0.0005);
            and children who showed postnatal catch-up growth (i.e. those who showed gains in
            weight or length between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other
			children (P=0.02; http://univ.edu.es/study.html) [20-22].

			Independent of current body composition, IGF-I levels at 5 yr were significantly
            associated with rate of weight gain between 0-2 yr (beta=0.19; P<0.0005);
            and children who showed postnatal catch-up growth (i.e. those who showed gains in
            weight or length between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other
			children (P=0.02; http://univ.edu.es/study.html) [20-22].

			Independent of current body composition, IGF-I levels at 5 yr were significantly
            associated with rate of weight gain between 0-2 yr (beta=0.19; P<0.0005);
            and children who showed postnatal catch-up growth (i.e. those who showed gains in
            weight or length between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other
			children (P=0.02; http://univ.edu.es/study.html) [20-22].`

	expected := []string{"Independent", "of", "current", "body", "composition", ",", "IGF-I",
		"levels", "at", "5", "yr", "were", "significantly", "associated", "with", "rate", "of",
		"weight", "gain", "between", "0-2", "yr", "(", "beta=0.19", ";", "P<0.0005", ")", ";",
		"and", "children", "who", "showed", "postnatal", "catch-up", "growth", "(", "i.e.", "those",
		"who", "showed", "gains", "in", "weight", "or", "length", "between", "0-2", "yr", "by",
		">0.67", "SD", "score", ")", "had", "higher", "IGF-I", "levels", "than", "other", "children",
		"(", "P=0.02", ";", "http://univ.edu.es/study.html", ")", "[", "20-22", "]", ".", "Independent", "of", "current", "body", "composition", ",", "IGF-I",
		"levels", "at", "5", "yr", "were", "significantly", "associated", "with", "rate", "of",
		"weight", "gain", "between", "0-2", "yr", "(", "beta=0.19", ";", "P<0.0005", ")", ";",
		"and", "children", "who", "showed", "postnatal", "catch-up", "growth", "(", "i.e.", "those",
		"who", "showed", "gains", "in", "weight", "or", "length", "between", "0-2", "yr", "by",
		">0.67", "SD", "score", ")", "had", "higher", "IGF-I", "levels", "than", "other", "children",
		"(", "P=0.02", ";", "http://univ.edu.es/study.html", ")", "[", "20-22", "]", ".", "Independent", "of", "current", "body", "composition", ",", "IGF-I",
		"levels", "at", "5", "yr", "were", "significantly", "associated", "with", "rate", "of",
		"weight", "gain", "between", "0-2", "yr", "(", "beta=0.19", ";", "P<0.0005", ")", ";",
		"and", "children", "who", "showed", "postnatal", "catch-up", "growth", "(", "i.e.", "those",
		"who", "showed", "gains", "in", "weight", "or", "length", "between", "0-2", "yr", "by",
		">0.67", "SD", "score", ")", "had", "higher", "IGF-I", "levels", "than", "other", "children",
		"(", "P=0.02", ";", "http://univ.edu.es/study.html", ")", "[", "20-22", "]", "."}

	doc, _ := makeDoc(web)
	assert.Equal(t, expected, getTokenText(doc))
}

func TestTokenizationTwitter(t *testing.T) {
	doc, _ := makeDoc("@twitter, what time does it start :-)")
	expected := []string{"@twitter", ",", "what", "time", "does", "it", "start", ":-)"}
	assert.Equal(t, expected, getTokenText(doc))

	doc, _ = makeDoc("Mr. James plays basketball in the N.B.A., do you?")
	expected = []string{
		"Mr.", "James", "plays", "basketball", "in", "the", "N.B.A.", ",",
		"do", "you", "?"}
	assert.Equal(t, expected, getTokenText(doc))
}

func BenchmarkTokenization(b *testing.B) {
	for n := 0; n < b.N; n++ {
		for _, s := range getWordBenchData() {
			makeDoc(s)
		}
	}
}
<commit_msg>Update tokenization benchmark<commit_after>package prose

import (
	"encoding/json"
	"path/filepath"
	"testing"

	"github.com/stretchr/testify/assert"
)

var testdata = "testdata"

func makeDoc(text string) (*Document, error) {
	return NewDocument(
		text,
		WithSegmentation(false),
		WithTagging(false),
		WithExtraction(false))
}

func getTokenText(doc *Document) []string {
	observed := []string{}
	for _, tok := range doc.Tokens() {
		observed = append(observed, tok.Text)
	}
	return observed
}

func getWordData(file string) ([]string, [][]string) {
	in := readDataFile(filepath.Join(testdata, "treebank_sents.json"))
	out := readDataFile(filepath.Join(testdata, file))

	input := []string{}
	output := [][]string{}

	checkError(json.Unmarshal(in, &input))
	checkError(json.Unmarshal(out, &output))

	return input, output
}

func getWordBenchData() []string {
	in := readDataFile(filepath.Join(testdata, "treebank_sents.json"))
	input := []string{}
	checkError(json.Unmarshal(in, &input))
	return input
}

func TestTokenizationSimple(t *testing.T) {
	doc, _ := makeDoc("Vale is a natural language linter that supports plain text, markup (Markdown, reStructuredText, AsciiDoc, and HTML), and source code comments. Vale doesn't attempt to offer a one-size-fits-all collection of rules—instead, it strives to make customization as easy as possible.")
	expected := []string{
		"Vale", "is", "a", "natural", "language", "linter", "that", "supports",
		"plain", "text", ",", "markup", "(", "Markdown", ",", "reStructuredText",
		",", "AsciiDoc", ",", "and", "HTML", ")", ",", "and", "source",
		"code", "comments", ".", "Vale", "does", "n't", "attempt", "to",
		"offer", "a", "one-size-fits-all", "collection", "of", "rules—instead",
		",", "it", "strives", "to", "make", "customization", "as", "easy", "as",
		"possible", "."}
	assert.Equal(t, expected, getTokenText(doc))
}

func TestTokenizationTreebank(t *testing.T) {
	input, output := getWordData("treebank_words.json")
	for i, s := range input {
		doc, _ := makeDoc(s)
		assert.Equal(t, output[i], getTokenText(doc))
	}
}

func TestTokenizationWeb(t *testing.T) {
	web := `Independent of current body composition, IGF-I levels at 5 yr were significantly
            associated with rate of weight gain between 0-2 yr (beta=0.19; P<0.0005);
            and children who showed postnatal catch-up growth (i.e. those who showed gains in
            weight or length between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other
			children (P=0.02; http://univ.edu.es/study.html) [20-22].`
	expected := []string{"Independent", "of", "current", "body", "composition", ",", "IGF-I",
		"levels", "at", "5", "yr", "were", "significantly", "associated", "with", "rate", "of",
		"weight", "gain", "between", "0-2", "yr", "(", "beta=0.19", ";", "P<0.0005", ")", ";",
		"and", "children", "who", "showed", "postnatal", "catch-up", "growth", "(", "i.e.", "those",
		"who", "showed", "gains", "in", "weight", "or", "length", "between", "0-2", "yr", "by",
		">0.67", "SD", "score", ")", "had", "higher", "IGF-I", "levels", "than", "other", "children",
		"(", "P=0.02", ";", "http://univ.edu.es/study.html", ")", "[", "20-22", "]", "."}
	doc, _ := makeDoc(web)
	assert.Equal(t, expected, getTokenText(doc))
}

func TestTokenizationWebParagraph(t *testing.T) {
	web := `Independent of current body composition, IGF-I levels at 5 yr were significantly
            associated with rate of weight gain between 0-2 yr (beta=0.19; P<0.0005);
            and children who showed postnatal catch-up growth (i.e. those who showed gains in
            weight or length between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other
			children (P=0.02; http://univ.edu.es/study.html) [20-22].

			Independent of current body composition, IGF-I levels at 5 yr were significantly
            associated with rate of weight gain between 0-2 yr (beta=0.19; P<0.0005);
            and children who showed postnatal catch-up growth (i.e. those who showed gains in
            weight or length between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other
			children (P=0.02; http://univ.edu.es/study.html) [20-22].

			Independent of current body composition, IGF-I levels at 5 yr were significantly
            associated with rate of weight gain between 0-2 yr (beta=0.19; P<0.0005);
            and children who showed postnatal catch-up growth (i.e. those who showed gains in
            weight or length between 0-2 yr by >0.67 SD score) had higher IGF-I levels than other
			children (P=0.02; http://univ.edu.es/study.html) [20-22].`

	expected := []string{"Independent", "of", "current", "body", "composition", ",", "IGF-I",
		"levels", "at", "5", "yr", "were", "significantly", "associated", "with", "rate", "of",
		"weight", "gain", "between", "0-2", "yr", "(", "beta=0.19", ";", "P<0.0005", ")", ";",
		"and", "children", "who", "showed", "postnatal", "catch-up", "growth", "(", "i.e.", "those",
		"who", "showed", "gains", "in", "weight", "or", "length", "between", "0-2", "yr", "by",
		">0.67", "SD", "score", ")", "had", "higher", "IGF-I", "levels", "than", "other", "children",
		"(", "P=0.02", ";", "http://univ.edu.es/study.html", ")", "[", "20-22", "]", ".", "Independent", "of", "current", "body", "composition", ",", "IGF-I",
		"levels", "at", "5", "yr", "were", "significantly", "associated", "with", "rate", "of",
		"weight", "gain", "between", "0-2", "yr", "(", "beta=0.19", ";", "P<0.0005", ")", ";",
		"and", "children", "who", "showed", "postnatal", "catch-up", "growth", "(", "i.e.", "those",
		"who", "showed", "gains", "in", "weight", "or", "length", "between", "0-2", "yr", "by",
		">0.67", "SD", "score", ")", "had", "higher", "IGF-I", "levels", "than", "other", "children",
		"(", "P=0.02", ";", "http://univ.edu.es/study.html", ")", "[", "20-22", "]", ".", "Independent", "of", "current", "body", "composition", ",", "IGF-I",
		"levels", "at", "5", "yr", "were", "significantly", "associated", "with", "rate", "of",
		"weight", "gain", "between", "0-2", "yr", "(", "beta=0.19", ";", "P<0.0005", ")", ";",
		"and", "children", "who", "showed", "postnatal", "catch-up", "growth", "(", "i.e.", "those",
		"who", "showed", "gains", "in", "weight", "or", "length", "between", "0-2", "yr", "by",
		">0.67", "SD", "score", ")", "had", "higher", "IGF-I", "levels", "than", "other", "children",
		"(", "P=0.02", ";", "http://univ.edu.es/study.html", ")", "[", "20-22", "]", "."}

	doc, _ := makeDoc(web)
	assert.Equal(t, expected, getTokenText(doc))
}

func TestTokenizationTwitter(t *testing.T) {
	doc, _ := makeDoc("@twitter, what time does it start :-)")
	expected := []string{"@twitter", ",", "what", "time", "does", "it", "start", ":-)"}
	assert.Equal(t, expected, getTokenText(doc))

	doc, _ = makeDoc("Mr. James plays basketball in the N.B.A., do you?")
	expected = []string{
		"Mr.", "James", "plays", "basketball", "in", "the", "N.B.A.", ",",
		"do", "you", "?"}
	assert.Equal(t, expected, getTokenText(doc))
}

func BenchmarkTokenization(b *testing.B) {
	in := readDataFile(filepath.Join(testdata, "sherlock.txt"))
	text := string(in)
	for n := 0; n < b.N; n++ {
		makeDoc(text)
	}
}
<|endoftext|>