<commit_before>// Copyright 2012-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.

// legacy_boot allows to handover a system running linuxboot/u-root
// to a legacy preinstalled operating system by replacing the traditional
// bootloader path

//
// Synopsis:
//	legacy_boot
//
// Description:
//	If returns to u-root shell, the code didn't found a local bootable option
//
// Notes:
//	The code is looking for boot/grub/grub.cfg file as to identify the
//	boot option.
//	The first bootable device found in the block device tree is the one used
//	Windows is not supported (that is a work in progress)
//
// Example:
//	legacy_boot -v 	- Start the script in verbose mode for debugging purpose

package main

import (
	"bufio"
	"fmt"
	"github.com/u-root/u-root/pkg/kexec"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"flag"
	"syscall"
)

var verbose bool
type options struct {
        verbose bool
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}
func check_fs(e error) {
}

func blkDevicesList(blkpath string, devpath string) []string {
	var blkDevices []string
	files, err := ioutil.ReadDir(blkpath)
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		deviceEntry, err := ioutil.ReadDir(blkpath + file.Name() + devpath)
		if err != nil {
			if ( verbose ) {
				println("can t read directory")
			}
			continue
		}
		blkDevices = append(blkDevices, deviceEntry[0].Name())
	}
	return (blkDevices)
}

// checkForBootableMbr is looking for bootable MBR signature
// Current support is limited to Hard disk devices and USB devices
func checkForBootableMbr(path string) int {
	var b511, b510 byte
	var err error
	f, err := os.Open(path)
	check(err)
	b1 := make([]byte, 512)
	_,err=f.Read(b1)
	check(err)
	b511 = b1[511]
	b510 = b1[510]
	err=f.Close()
	check(err)
	if ((b511 == 0xaa) && (b510 == 0x55)) == true {
		return 1
	}
	return 0
}

// getDevicePartList returns all devices attached to a specific name like /dev/sdaX where X can move from 0 to 127
// FIXME no support for devices which are included into subdirectory within /dev
func getDevicePartList(path string) []string {
	var returnValue []string
	files, err := ioutil.ReadDir("/dev/")
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		if strings.HasPrefix(file.Name(), path) {
			// We shall not return full device name
			if file.Name() != path {
				// We shall check that the remaining part is a number
				returnValue = append(returnValue, file.Name())
			}
		}
	}
	return returnValue
}

// getSupportedFilesystem returns all block file system supported by the linuxboot kernel
func getSupportedFilesystem() []string {
	var returnValue []string
	var err error
	file, err := os.Open("/proc/filesystems")
	check(err)
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		log.Fatal(err)
	}

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if fields[0] == "nodev" {
			continue
		}
		if fields[0] != "" {
			returnValue = append(returnValue, fields[0])
    		}
	}

	err=file.Close()
	check(err)
	return returnValue

}

// mountEntry tries to mount a specific block device
func mountEntry(path string, supportedFilesystem []string) bool {
	var returnValue bool
	var err error
	err=syscall.Mkdir("/u-root", 0777)
	check_fs(err)
	var flags uintptr
	// Was supposed to be unecessary for kernel 4.x.x
	if verbose {
		println("/dev/" + path)
	}
	for _, filesystem := range supportedFilesystem {
		flags = syscall.MS_MGC_VAL
		// Need to load the filesystem kind supported
		err = syscall.Mkdir("/u-root/"+path, 0777)
		check_fs(err)
		err := syscall.Mount("/dev/"+path, "/u-root/"+path, filesystem, flags, "")
		if err == nil {
			return true
		}
	}
	returnValue = false
	return returnValue
}

func umountEntry(path string) bool {
	var returnValue bool
	var flags int
	// Was supposed to be unecessary for kernel 4.x.x
	flags = syscall.MNT_DETACH
	err := syscall.Unmount(path, flags)
	if err == nil {
		return true
	}
	returnValue = false
	return returnValue
}

// checkBootEntry is looking for grub.cfg file
// and return absolute path to it
func checkBootEntry(mountPoint string) string {
	_, err := os.Stat(mountPoint + "/boot")
	if err == nil {
		// The boot directory is there
		_, err2 := os.Stat(mountPoint + "/boot" + "/grub")
		if err2 == nil {
			_, err3 := os.Stat(mountPoint + "/boot" + "/grub" + "/grub.cfg")
			if err3 == nil {
				// found
				return (mountPoint + "/boot" + "/grub")
			}
		}
	}
	return ""

}

// getFileMenuContent is parsing a grub.cfg file
// input: absolute directory path to grub.cfg
// output: Return a list of strings with the following format
//	 line[3*x] - menuconfig
//	 line[3*x+1] - linux kernel + boot options
// 	 line[3*x+2] - initrd
// and the default boot entry configured into grub.cfg
func getFileMenuContent(path string) ([]string, int) {
	var returnValue []string
	var err error
	file, err := os.Open(path + "/grub.cfg")
	check(err)
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		log.Fatal(err)
	}
	var status int
	var intReturn int
	intReturn = 0
	status = 0
	// When status = 0 we are looking for a menu entry
	// When status = 1 we are looking for a linux entry
	// When status = 2 we are looking for a initrd entry
	var trimmedLine string
	for scanner.Scan() {
		trimmedLine = strings.TrimSpace(scanner.Text())
		trimmedLine = strings.Join(strings.Fields(trimmedLine), " ")
		if !strings.HasPrefix(trimmedLine, "#") {
			if (strings.HasPrefix(trimmedLine, "set default=") ) && (status == 0) {
				fmt.Sscanf(trimmedLine, "set default=\"%d\"", &intReturn)
			}
			if (strings.HasPrefix(trimmedLine, "menuentry ")) && (status == 0) {
				status = 1
				returnValue = append(returnValue, trimmedLine)
			}
			if (strings.HasPrefix(trimmedLine, "linux ")) && (status == 1) {
				status = 2
				returnValue = append(returnValue, trimmedLine)
			}
			if (strings.HasPrefix(trimmedLine, "initrd ")) && (status == 2) {
				status = 0
				returnValue = append(returnValue, trimmedLine)
			}
		}
	}
	err = file.Close()
	check(err)
	return returnValue, intReturn

}

func copyLocal(path string) string {
	var dest string
	var err error
	result := strings.Split(path, "/")
	for _, entry := range result {
		dest = entry
	}
	dest = "/tmp/" + dest
	srcFile, err := os.Open(path)
	check(err)

	destFile, err := os.Create(dest) // creates if file doesn't exist
	check(err)

	_, err = io.Copy(destFile, srcFile) // check first var for number of bytes copied
	check(err)

	err = destFile.Sync()
	check(err)
	err = destFile.Close()
	check(err)
	err = srcFile.Close()
	check(err)
	return dest
}

// kexecEntry is booting new kernel based on the content of grub.cfg
func kexecEntry(grubConfPath string, mountPoint string) {
	var fileMenuContent []string
	var entry int
	var localKernelPath string
	var localInitrdPath string
	if verbose {
		println(grubConfPath)
	}
	fileMenuContent, entry = getFileMenuContent(grubConfPath)
	var kernel string
	var kernelParameter string
	var initrd string
	var kernelInfos []string
	kernelInfos = strings.Fields(fileMenuContent[3*entry+1])
	kernel = kernelInfos[1]
	var count int
	count = 0
	for _, field := range kernelInfos {
		if count > 1 {
			kernelParameter = kernelParameter + " " + field
		}
		count = count + 1
	}
	fmt.Sscanf(fileMenuContent[3*entry+2], "initrd %s", &initrd)
	if verbose {
		println("************** boot parameters  ********************")
		println(kernel)
		println(kernelParameter)
		println(initrd)
		println("****************************************************")
	}
	localKernelPath = copyLocal(mountPoint + kernel)
	localInitrdPath = copyLocal(mountPoint + initrd)
	if verbose {
		println(localKernelPath)
	}
	umountEntry(mountPoint)
	// We can kexec the kernel with localKernelPath as kernel entry, kernelParameter as parameter and initrd as initrd !
	log.Printf("Loading %s for kernel\n", localKernelPath)

	kernelDesc, err := os.OpenFile(localKernelPath, os.O_RDONLY, 0)
	if err != nil {
		log.Fatalf("open(%q): %v", localKernelPath, err)
	}
	// defer kernelDesc.Close()

	var ramfs *os.File
	ramfs, err = os.OpenFile(localInitrdPath, os.O_RDONLY, 0)
	if err != nil {
		log.Fatalf("open(%q): %v", localInitrdPath, err)
	}
	// defer ramfs.Close()

	if err := kexec.FileLoad(kernelDesc, ramfs, kernelParameter); err != nil {
		log.Fatalf("%v", err)
	}
 	err = ramfs.Close()
	check(err)
	err = kernelDesc.Close()
	check(err)
	if err := kexec.Reboot(); err != nil {
		log.Fatalf("%v", err)
	}

}

func registerFlags(f *flag.FlagSet) *options {
        o := &options{}
        f.BoolVar(&o.verbose, "v", false, "Set verbose output")
        return o
}

func main() {
	var blkList []string
	var supportedFilesystem []string
	verbose = false
        opts := registerFlags(flag.CommandLine)
	flag.Parse()

        if opts.verbose != false {
		verbose = true
        }

	supportedFilesystem = getSupportedFilesystem()
	if verbose {
		println("************** Supported Filesystem by current linuxboot ********************")
		for _, filesystem := range supportedFilesystem {
			println(filesystem)
		}
		println("*****************************************************************************")
	}
	blkList = blkDevicesList("/sys/dev/block/", "/device/block/")
	// We must validate if the MBR is bootable or not and keep the
	// devices which do have such support
	// drive are easy to detect
	for _, entry := range blkList {
		if checkForBootableMbr("/dev/"+entry) == 1 {
			fmt.Println("Bootable device found")
			// We need to loop on the device entries which are into /dev/<device>X
			// and mount each partitions as to find /boot entry if it is available somewhere
			var devicePartList []string
			devicePartList = getDevicePartList(entry)
			for _, deviceList := range devicePartList {
				if mountEntry(deviceList, supportedFilesystem) {
					if verbose {
						println("mount succeed")
					}
					var grubConfPath = checkBootEntry("/u-root/" + deviceList)
					if grubConfPath != "" {
						if verbose {
							println("calling basic kexec")
						}
						kexecEntry(grubConfPath, "/u-root/"+deviceList)
					}
				}
				umountEntry("/u-root/" + deviceList)
			}
		}
	}
	println("Sorry no bootable device found")
}
<commit_msg>Add Ron suggestion on improvement regarding code style to detect bootable MBR signature<commit_after>// Copyright 2012-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.

// legacy_boot allows to handover a system running linuxboot/u-root
// to a legacy preinstalled operating system by replacing the traditional
// bootloader path

//
// Synopsis:
//	legacy_boot
//
// Description:
//	If returns to u-root shell, the code didn't found a local bootable option
//
// Notes:
//	The code is looking for boot/grub/grub.cfg file as to identify the
//	boot option.
//	The first bootable device found in the block device tree is the one used
//	Windows is not supported (that is a work in progress)
//
// Example:
//	legacy_boot -v 	- Start the script in verbose mode for debugging purpose

package main

import (
	"bufio"
	"encoding/binary"
	"fmt"
	"github.com/u-root/u-root/pkg/kexec"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"flag"
	"syscall"
)

const (
	bootableMBR = 0xaa55
	signatureOffset = 510
)

var verbose bool
type options struct {
        verbose bool
}

func check(e error) {
	if e != nil {
		panic(e)
	}
}
func check_fs(e error) {
}

func blkDevicesList(blkpath string, devpath string) []string {
	var blkDevices []string
	files, err := ioutil.ReadDir(blkpath)
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		deviceEntry, err := ioutil.ReadDir(blkpath + file.Name() + devpath)
		if err != nil {
			if ( verbose ) {
				println("can t read directory")
			}
			continue
		}
		blkDevices = append(blkDevices, deviceEntry[0].Name())
	}
	return (blkDevices)
}

// checkForBootableMBR is looking for bootable MBR signature
// Current support is limited to Hard disk devices and USB devices
func checkForBootableMBR(path string) (bool,error) {
	var sig uint16
	f, err := os.Open(path)
	if err != nil {
		return false, err
	}
	if err := binary.Read(io.NewSectionReader(f, signatureOffset, 2), binary.LittleEndian, &sig); err != nil {
		return false,err
	}
	if sig != bootableMBR {
		return false, nil
	}
	return true,nil
}

// getDevicePartList returns all devices attached to a specific name like /dev/sdaX where X can move from 0 to 127
// FIXME no support for devices which are included into subdirectory within /dev
func getDevicePartList(path string) []string {
	var returnValue []string
	files, err := ioutil.ReadDir("/dev/")
	if err != nil {
		log.Fatal(err)
	}
	for _, file := range files {
		if strings.HasPrefix(file.Name(), path) {
			// We shall not return full device name
			if file.Name() != path {
				// We shall check that the remaining part is a number
				returnValue = append(returnValue, file.Name())
			}
		}
	}
	return returnValue
}

// getSupportedFilesystem returns all block file system supported by the linuxboot kernel
func getSupportedFilesystem() []string {
	var returnValue []string
	var err error
	file, err := os.Open("/proc/filesystems")
	check(err)
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		log.Fatal(err)
	}

	for scanner.Scan() {
		fields := strings.Fields(scanner.Text())
		if fields[0] == "nodev" {
			continue
		}
		if fields[0] != "" {
			returnValue = append(returnValue, fields[0])
    		}
	}

	err=file.Close()
	check(err)
	return returnValue

}

// mountEntry tries to mount a specific block device
func mountEntry(path string, supportedFilesystem []string) bool {
	var returnValue bool
	var err error
	err=syscall.Mkdir("/u-root", 0777)
	check_fs(err)
	var flags uintptr
	// Was supposed to be unecessary for kernel 4.x.x
	if verbose {
		println("/dev/" + path)
	}
	for _, filesystem := range supportedFilesystem {
		flags = syscall.MS_MGC_VAL
		// Need to load the filesystem kind supported
		err = syscall.Mkdir("/u-root/"+path, 0777)
		check_fs(err)
		err := syscall.Mount("/dev/"+path, "/u-root/"+path, filesystem, flags, "")
		if err == nil {
			return true
		}
	}
	returnValue = false
	return returnValue
}

func umountEntry(path string) bool {
	var returnValue bool
	var flags int
	// Was supposed to be unecessary for kernel 4.x.x
	flags = syscall.MNT_DETACH
	err := syscall.Unmount(path, flags)
	if err == nil {
		return true
	}
	returnValue = false
	return returnValue
}

// checkBootEntry is looking for grub.cfg file
// and return absolute path to it
func checkBootEntry(mountPoint string) string {
	_, err := os.Stat(mountPoint + "/boot")
	if err == nil {
		// The boot directory is there
		_, err2 := os.Stat(mountPoint + "/boot" + "/grub")
		if err2 == nil {
			_, err3 := os.Stat(mountPoint + "/boot" + "/grub" + "/grub.cfg")
			if err3 == nil {
				// found
				return (mountPoint + "/boot" + "/grub")
			}
		}
	}
	return ""

}

// getFileMenuContent is parsing a grub.cfg file
// input: absolute directory path to grub.cfg
// output: Return a list of strings with the following format
//	 line[3*x] - menuconfig
//	 line[3*x+1] - linux kernel + boot options
// 	 line[3*x+2] - initrd
// and the default boot entry configured into grub.cfg
func getFileMenuContent(path string) ([]string, int) {
	var returnValue []string
	var err error
	file, err := os.Open(path + "/grub.cfg")
	check(err)
	scanner := bufio.NewScanner(file)
	if err = scanner.Err(); err != nil {
		log.Fatal(err)
	}
	var status int
	var intReturn int
	intReturn = 0
	status = 0
	// When status = 0 we are looking for a menu entry
	// When status = 1 we are looking for a linux entry
	// When status = 2 we are looking for a initrd entry
	var trimmedLine string
	for scanner.Scan() {
		trimmedLine = strings.TrimSpace(scanner.Text())
		trimmedLine = strings.Join(strings.Fields(trimmedLine), " ")
		if !strings.HasPrefix(trimmedLine, "#") {
			if (strings.HasPrefix(trimmedLine, "set default=") ) && (status == 0) {
				fmt.Sscanf(trimmedLine, "set default=\"%d\"", &intReturn)
			}
			if (strings.HasPrefix(trimmedLine, "menuentry ")) && (status == 0) {
				status = 1
				returnValue = append(returnValue, trimmedLine)
			}
			if (strings.HasPrefix(trimmedLine, "linux ")) && (status == 1) {
				status = 2
				returnValue = append(returnValue, trimmedLine)
			}
			if (strings.HasPrefix(trimmedLine, "initrd ")) && (status == 2) {
				status = 0
				returnValue = append(returnValue, trimmedLine)
			}
		}
	}
	err = file.Close()
	check(err)
	return returnValue, intReturn

}

func copyLocal(path string) string {
	var dest string
	var err error
	result := strings.Split(path, "/")
	for _, entry := range result {
		dest = entry
	}
	dest = "/tmp/" + dest
	srcFile, err := os.Open(path)
	check(err)

	destFile, err := os.Create(dest) // creates if file doesn't exist
	check(err)

	_, err = io.Copy(destFile, srcFile) // check first var for number of bytes copied
	check(err)

	err = destFile.Sync()
	check(err)
	err = destFile.Close()
	check(err)
	err = srcFile.Close()
	check(err)
	return dest
}

// kexecEntry is booting new kernel based on the content of grub.cfg
func kexecEntry(grubConfPath string, mountPoint string) {
	var fileMenuContent []string
	var entry int
	var localKernelPath string
	var localInitrdPath string
	if verbose {
		println(grubConfPath)
	}
	fileMenuContent, entry = getFileMenuContent(grubConfPath)
	var kernel string
	var kernelParameter string
	var initrd string
	var kernelInfos []string
	kernelInfos = strings.Fields(fileMenuContent[3*entry+1])
	kernel = kernelInfos[1]
	var count int
	count = 0
	for _, field := range kernelInfos {
		if count > 1 {
			kernelParameter = kernelParameter + " " + field
		}
		count = count + 1
	}
	fmt.Sscanf(fileMenuContent[3*entry+2], "initrd %s", &initrd)
	if verbose {
		println("************** boot parameters  ********************")
		println(kernel)
		println(kernelParameter)
		println(initrd)
		println("****************************************************")
	}
	localKernelPath = copyLocal(mountPoint + kernel)
	localInitrdPath = copyLocal(mountPoint + initrd)
	if verbose {
		println(localKernelPath)
	}
	umountEntry(mountPoint)
	// We can kexec the kernel with localKernelPath as kernel entry, kernelParameter as parameter and initrd as initrd !
	log.Printf("Loading %s for kernel\n", localKernelPath)

	kernelDesc, err := os.OpenFile(localKernelPath, os.O_RDONLY, 0)
	if err != nil {
		log.Fatalf("open(%q): %v", localKernelPath, err)
	}
	// defer kernelDesc.Close()

	var ramfs *os.File
	ramfs, err = os.OpenFile(localInitrdPath, os.O_RDONLY, 0)
	if err != nil {
		log.Fatalf("open(%q): %v", localInitrdPath, err)
	}
	// defer ramfs.Close()

	if err := kexec.FileLoad(kernelDesc, ramfs, kernelParameter); err != nil {
		log.Fatalf("%v", err)
	}
 	err = ramfs.Close()
	check(err)
	err = kernelDesc.Close()
	check(err)
	if err := kexec.Reboot(); err != nil {
		log.Fatalf("%v", err)
	}

}

func registerFlags(f *flag.FlagSet) *options {
        o := &options{}
        f.BoolVar(&o.verbose, "v", false, "Set verbose output")
        return o
}

func main() {
	var blkList []string
	var supportedFilesystem []string
	verbose = false
        opts := registerFlags(flag.CommandLine)
	flag.Parse()

        if opts.verbose != false {
		verbose = true
        }

	supportedFilesystem = getSupportedFilesystem()
	if verbose {
		println("************** Supported Filesystem by current linuxboot ********************")
		for _, filesystem := range supportedFilesystem {
			println(filesystem)
		}
		println("*****************************************************************************")
	}
	blkList = blkDevicesList("/sys/dev/block/", "/device/block/")
	// We must validate if the MBR is bootable or not and keep the
	// devices which do have such support
	// drive are easy to detect
	for _, entry := range blkList {
		dev := filepath.Join("/dev", entry)
		b, err := checkForBootableMBR(dev)
		if err != nil {
			// Not sure it matters; there can be many bogus entries?
			log.Printf("MBR for %s failed: %v", dev, err)
		}
		if b {
			fmt.Println("Bootable device found")
			// We need to loop on the device entries which are into /dev/<device>X
			// and mount each partitions as to find /boot entry if it is available somewhere
			var devicePartList []string
			devicePartList = getDevicePartList(entry)
			for _, deviceList := range devicePartList {
				if mountEntry(deviceList, supportedFilesystem) {
					if verbose {
						println("mount succeed")
					}
					var grubConfPath = checkBootEntry("/u-root/" + deviceList)
					if grubConfPath != "" {
						if verbose {
							println("calling basic kexec")
						}
						kexecEntry(grubConfPath, "/u-root/"+deviceList)
					}
				}
				umountEntry("/u-root/" + deviceList)
			}
		}
	}
	println("Sorry no bootable device found")
}
<|endoftext|>


<commit_before>package drainer

import (
	"context"
	"testing"
	"time"

	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/state"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"
	"github.com/stretchr/testify/require"
	"golang.org/x/time/rate"
)

func testNodeDrainWatcher(t *testing.T) (*nodeDrainWatcher, *state.StateStore, *MockNodeTracker) {
	t.Helper()
	state := state.TestStateStore(t)
	limiter := rate.NewLimiter(100.0, 100)
	logger := testlog.HCLogger(t)
	m := NewMockNodeTracker()
	w := NewNodeDrainWatcher(context.Background(), limiter, state, logger, m)
	return w, state, m
}

func TestNodeDrainWatcher_Interface(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	w, _, _ := testNodeDrainWatcher(t)
	require.Implements((*DrainingNodeWatcher)(nil), w)
}

func TestNodeDrainWatcher_AddDraining(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, state, m := testNodeDrainWatcher(t)

	// Create two nodes, one draining and one not draining
	n1, n2 := mock.Node(), mock.Node()
	n2.DrainStrategy = &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: time.Hour,
		},
		ForceDeadline: time.Now().Add(time.Hour),
	}

	require.Nil(state.UpsertNode(100, n1))
	require.Nil(state.UpsertNode(101, n2))

	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 1, nil
	}, func(err error) {
		t.Fatal("No node drain events")
	})

	tracked := m.TrackedNodes()
	require.NotContains(tracked, n1.ID)
	require.Contains(tracked, n2.ID)
	require.Equal(n2, tracked[n2.ID])

}

func TestNodeDrainWatcher_Remove(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, state, m := testNodeDrainWatcher(t)

	// Create a draining node
	n := mock.Node()
	n.DrainStrategy = &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: time.Hour,
		},
		ForceDeadline: time.Now().Add(time.Hour),
	}

	// Wait for it to be tracked
	require.Nil(state.UpsertNode(100, n))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 1, nil
	}, func(err error) {
		t.Fatal("No node drain events")
	})

	tracked := m.TrackedNodes()
	require.Contains(tracked, n.ID)
	require.Equal(n, tracked[n.ID])

	// Change the node to be not draining and wait for it to be untracked
	require.Nil(state.UpdateNodeDrain(101, n.ID, nil, false, 0, nil))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 2, nil
	}, func(err error) {
		t.Fatal("No new node drain events")
	})

	tracked = m.TrackedNodes()
	require.NotContains(tracked, n.ID)
}

func TestNodeDrainWatcher_Remove_Nonexistent(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, state, m := testNodeDrainWatcher(t)

	// Create a draining node
	n := mock.Node()
	n.DrainStrategy = &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: time.Hour,
		},
		ForceDeadline: time.Now().Add(time.Hour),
	}

	// Wait for it to be tracked
	require.Nil(state.UpsertNode(100, n))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 1, nil
	}, func(err error) {
		t.Fatal("No node drain events")
	})

	tracked := m.TrackedNodes()
	require.Contains(tracked, n.ID)
	require.Equal(n, tracked[n.ID])

	// Delete the node
	require.Nil(state.DeleteNode(101, n.ID))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 2, nil
	}, func(err error) {
		t.Fatal("No new node drain events")
	})

	tracked = m.TrackedNodes()
	require.NotContains(tracked, n.ID)
}

func TestNodeDrainWatcher_Update(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, state, m := testNodeDrainWatcher(t)

	// Create a draining node
	n := mock.Node()
	n.DrainStrategy = &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: time.Hour,
		},
		ForceDeadline: time.Now().Add(time.Hour),
	}

	// Wait for it to be tracked
	require.Nil(state.UpsertNode(100, n))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 1, nil
	}, func(err error) {
		t.Fatal("No node drain events")
	})

	tracked := m.TrackedNodes()
	require.Contains(tracked, n.ID)
	require.Equal(n, tracked[n.ID])

	// Change the node to have a new spec
	s2 := n.DrainStrategy.Copy()
	s2.Deadline += time.Hour
	require.Nil(state.UpdateNodeDrain(101, n.ID, s2, false, 0, nil))

	// Wait for it to be updated
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 2, nil
	}, func(err error) {
		t.Fatal("No new node drain events")
	})

	tracked = m.TrackedNodes()
	require.Contains(tracked, n.ID)
	require.Equal(s2, tracked[n.ID].DrainStrategy)
}
<commit_msg>drainer watch_nodes_test batch of 1<commit_after>package drainer

import (
	"context"
	"testing"
	"time"

	"github.com/hashicorp/nomad/helper/testlog"
	"github.com/hashicorp/nomad/nomad/mock"
	"github.com/hashicorp/nomad/nomad/state"
	"github.com/hashicorp/nomad/nomad/structs"
	"github.com/hashicorp/nomad/testutil"
	"github.com/stretchr/testify/require"
	"golang.org/x/time/rate"
)

func testNodeDrainWatcher(t *testing.T) (*nodeDrainWatcher, *state.StateStore, *MockNodeTracker) {
	t.Helper()
	state := state.TestStateStore(t)
	limiter := rate.NewLimiter(100.0, 100)
	logger := testlog.HCLogger(t)
	m := NewMockNodeTracker()
	w := NewNodeDrainWatcher(context.Background(), limiter, state, logger, m)
	return w, state, m
}

func TestNodeDrainWatcher_Interface(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	w, _, _ := testNodeDrainWatcher(t)
	require.Implements((*DrainingNodeWatcher)(nil), w)
}

func TestNodeDrainWatcher_AddDraining(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, state, m := testNodeDrainWatcher(t)

	// Create two nodes, one draining and one not draining
	n1, n2 := mock.Node(), mock.Node()
	n2.DrainStrategy = &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: time.Hour,
		},
		ForceDeadline: time.Now().Add(time.Hour),
	}

	require.Nil(state.UpsertNode(100, n1))
	require.Nil(state.UpsertNode(101, n2))

	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 1, nil
	}, func(err error) {
		t.Fatal("No node drain events")
	})

	tracked := m.TrackedNodes()
	require.NotContains(tracked, n1.ID)
	require.Contains(tracked, n2.ID)
	require.Equal(n2, tracked[n2.ID])

}

func TestNodeDrainWatcher_Remove(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, state, m := testNodeDrainWatcher(t)

	// Create a draining node
	n := mock.Node()
	n.DrainStrategy = &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: time.Hour,
		},
		ForceDeadline: time.Now().Add(time.Hour),
	}

	// Wait for it to be tracked
	require.Nil(state.UpsertNode(100, n))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 1, nil
	}, func(err error) {
		t.Fatal("No node drain events")
	})

	tracked := m.TrackedNodes()
	require.Contains(tracked, n.ID)
	require.Equal(n, tracked[n.ID])

	// Change the node to be not draining and wait for it to be untracked
	require.Nil(state.UpdateNodeDrain(101, n.ID, nil, false, 0, nil))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 2, nil
	}, func(err error) {
		t.Fatal("No new node drain events")
	})

	tracked = m.TrackedNodes()
	require.NotContains(tracked, n.ID)
}

func TestNodeDrainWatcher_Remove_Nonexistent(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, state, m := testNodeDrainWatcher(t)

	// Create a draining node
	n := mock.Node()
	n.DrainStrategy = &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: time.Hour,
		},
		ForceDeadline: time.Now().Add(time.Hour),
	}

	// Wait for it to be tracked
	require.Nil(state.UpsertNode(100, n))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 1, nil
	}, func(err error) {
		t.Fatal("No node drain events")
	})

	tracked := m.TrackedNodes()
	require.Contains(tracked, n.ID)
	require.Equal(n, tracked[n.ID])

	// Delete the node
	require.Nil(state.DeleteNode(101, []string{n.ID}))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 2, nil
	}, func(err error) {
		t.Fatal("No new node drain events")
	})

	tracked = m.TrackedNodes()
	require.NotContains(tracked, n.ID)
}

func TestNodeDrainWatcher_Update(t *testing.T) {
	t.Parallel()
	require := require.New(t)
	_, state, m := testNodeDrainWatcher(t)

	// Create a draining node
	n := mock.Node()
	n.DrainStrategy = &structs.DrainStrategy{
		DrainSpec: structs.DrainSpec{
			Deadline: time.Hour,
		},
		ForceDeadline: time.Now().Add(time.Hour),
	}

	// Wait for it to be tracked
	require.Nil(state.UpsertNode(100, n))
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 1, nil
	}, func(err error) {
		t.Fatal("No node drain events")
	})

	tracked := m.TrackedNodes()
	require.Contains(tracked, n.ID)
	require.Equal(n, tracked[n.ID])

	// Change the node to have a new spec
	s2 := n.DrainStrategy.Copy()
	s2.Deadline += time.Hour
	require.Nil(state.UpdateNodeDrain(101, n.ID, s2, false, 0, nil))

	// Wait for it to be updated
	testutil.WaitForResult(func() (bool, error) {
		return len(m.events()) == 2, nil
	}, func(err error) {
		t.Fatal("No new node drain events")
	})

	tracked = m.TrackedNodes()
	require.Contains(tracked, n.ID)
	require.Equal(s2, tracked[n.ID].DrainStrategy)
}
<|endoftext|>


<commit_before>package rigo

import (
	"fmt"
	"os"
	"time"

	. "github.com/mae-global/rigo/ri"
)

func DefaultFilePipe() *Pipe {
	pipe := NewPipe()
	return pipe.Append(&PipeTimer{}).Append(&PipeToStats{}).Append(&PipeToFile{})
}

/* Time from Begin to End */
type PipeTimer struct {
	start  time.Time
	finish time.Time
}

func (p PipeTimer) Name() string {
	return "default-pipe-timer"
}

func (p *PipeTimer) Write(name RtName, list []Rter, info Info) *Result {
	switch string(name) {
	case "Begin", "RiBegin":
		p.start = time.Now()
		p.finish = p.start
		break
	case "End", "RiEnd":
		p.finish = time.Now()
		break
	}
	return Done()
}

func (p *PipeTimer) String() string {
	return fmt.Sprintf("pipe took %s", p.finish.Sub(p.start))
}

func (p *PipeTimer) Took() time.Duration {
	return p.finish.Sub(p.start)
}

/* Pipe RI output to gathered states */
type PipeToStats struct {
	Stats map[RtName]int
}

func (p PipeToStats) Name() string {
	return "default-pipe-to-stats"
}

func (p *PipeToStats) Write(name RtName, list []Rter, info Info) *Result {
	if p.Stats == nil {
		p.Stats = make(map[RtName]int, 0)
	}
	if _, exists := p.Stats[name]; !exists {
		p.Stats[name] = 0
	}
	p.Stats[name]++

	return Done()
}

func (p *PipeToStats) String() string {
	if p.Stats == nil {
		return "stats [empty]"
	}

	if len(p.Stats) == 0 {
		return "stats [empty]"
	}

	max := 0
	for _, v := range p.Stats {
		if v > max {
			max = v
		}
	}

	dfmt := "\t%0" + fmt.Sprintf("%d", len(fmt.Sprintf("%d", max))) + "d"

	out := fmt.Sprintf("stats %d [\n", len(p.Stats))
	for n, v := range p.Stats {
		out += fmt.Sprintf(dfmt+" call(s).....%s\n", v, n)
	}
	return out + "]\n"
}

/* Pipe RI output to file */
type PipeToFile struct {
	file *os.File
}

func (p PipeToFile) Name() string {
	return "default-pipe-to-file"
}

func (p *PipeToFile) Write(name RtName, list []Rter, info Info) *Result {
	if info.Formal {
		name = name.Trim("Ri")
	}

	if name == "Begin" {
		if p.file != nil {
			return InError(ErrProtocolBotch)
		}
		file := "out.rib"
		if len(list) > 0 {
			if t, ok := list[0].(RtString); ok {
				file = string(t)
			}
		}

		f, err := os.Create(file)
		if err != nil {
			return InError(err)
		}
		p.file = f

		postfix := "\n"
		if info.Entity {
			postfix = " Entity\n"
		}
		if _, err = p.file.Write([]byte(string(RIBStructure) + postfix)); err != nil {
			return InError(err)
		}
		return Done()
	}

	if name == "End" {
		if p.file == nil {
			return InError(ErrProtocolBotch)
		}
		if err := p.file.Close(); err != nil {
			return InError(err)
		}
		return Done()
	}

	if p.file == nil {
		return InError(ErrNoActiveContext)
	}

	if name == "Verbatim" {
		if _, err := p.file.Write([]byte(Serialise(list) + "\n")); err != nil {
			return InError(err)
		}
		return Done()
	}

	if name != "##" {

		prefix := ""
		for i := 0; i < info.Depth; i++ {
			prefix += "\t"
		}

		if _, err := p.file.Write([]byte(prefix + name.Serialise() + " " + Serialise(list) + "\n")); err != nil {
			return InError(err)
		}
		return Done()
	}

	if _, err := p.file.Write([]byte("##" + Serialise(list) + "\n")); err != nil {
		return InError(err)
	}
	return Done()
}
<commit_msg>added sort to stats<commit_after>package rigo

import (
	"fmt"
	"os"
	"time"
	"sort"

	. "github.com/mae-global/rigo/ri"
)

func DefaultFilePipe() *Pipe {
	pipe := NewPipe()
	return pipe.Append(&PipeTimer{}).Append(&PipeToStats{}).Append(&PipeToFile{})
}

/* Time from Begin to End */
type PipeTimer struct {
	start  time.Time
	finish time.Time
}

func (p PipeTimer) Name() string {
	return "default-pipe-timer"
}

func (p *PipeTimer) Write(name RtName, list []Rter, info Info) *Result {
	switch string(name) {
	case "Begin", "RiBegin":
		p.start = time.Now()
		p.finish = p.start
		break
	case "End", "RiEnd":
		p.finish = time.Now()
		break
	}
	return Done()
}

func (p *PipeTimer) String() string {
	return fmt.Sprintf("pipe took %s", p.finish.Sub(p.start))
}

func (p *PipeTimer) Took() time.Duration {
	return p.finish.Sub(p.start)
}

/* Pipe RI output to gathered states */
type PipeToStats struct {
	Stats map[RtName]int
}

func (p PipeToStats) Name() string {
	return "default-pipe-to-stats"
}

func (p *PipeToStats) Write(name RtName, list []Rter, info Info) *Result {
	if p.Stats == nil {
		p.Stats = make(map[RtName]int, 0)
	}
	if _, exists := p.Stats[name]; !exists {
		p.Stats[name] = 0
	}
	p.Stats[name]++

	return Done()
}

type record struct {
	Name RtName
	Count int
}

type byCount []record

func (bc byCount) Len() int { return len(bc) }
func (bc byCount) Less(i,j int) bool { return bc[i].Count < bc[j].Count }
func (bc byCount) Swap(i,j int) { bc[i],bc[j] = bc[j],bc[i] }

func (p *PipeToStats) String() string {
	if p.Stats == nil {
		return "stats [empty]"
	}

	if len(p.Stats) == 0 {
		return "stats [empty]"
	}

	max := 0
	records := make([]record,0)

	for n, v := range p.Stats {
		if v > max {
			max = v
		}

		records = append(records,record{n,v})
	}

	sort.Sort(byCount(records))

	dfmt := "\t%0" + fmt.Sprintf("%d", len(fmt.Sprintf("%d", max))) + "d"

	out := fmt.Sprintf("stats %d [\n", len(p.Stats))
	for _, r := range records {
		out += fmt.Sprintf(dfmt + " call(s).....%s\n", r.Count,r.Name)
	}
	return out + "]\n"
}

/* Pipe RI output to file */
type PipeToFile struct {
	file *os.File
}

func (p PipeToFile) Name() string {
	return "default-pipe-to-file"
}

func (p *PipeToFile) Write(name RtName, list []Rter, info Info) *Result {
	if info.Formal {
		name = name.Trim("Ri")
	}

	if name == "Begin" {
		if p.file != nil {
			return InError(ErrProtocolBotch)
		}
		file := "out.rib"
		if len(list) > 0 {
			if t, ok := list[0].(RtString); ok {
				file = string(t)
			}
		}

		f, err := os.Create(file)
		if err != nil {
			return InError(err)
		}
		p.file = f

		postfix := "\n"
		if info.Entity {
			postfix = " Entity\n"
		}
		if _, err = p.file.Write([]byte(string(RIBStructure) + postfix)); err != nil {
			return InError(err)
		}
		return Done()
	}

	if name == "End" {
		if p.file == nil {
			return InError(ErrProtocolBotch)
		}
		if err := p.file.Close(); err != nil {
			return InError(err)
		}
		return Done()
	}

	if p.file == nil {
		return InError(ErrNoActiveContext)
	}

	if name == "Verbatim" {
		if _, err := p.file.Write([]byte(Serialise(list) + "\n")); err != nil {
			return InError(err)
		}
		return Done()
	}

	if name != "##" {

		prefix := ""
		for i := 0; i < info.Depth; i++ {
			prefix += "\t"
		}

		if _, err := p.file.Write([]byte(prefix + name.Serialise() + " " + Serialise(list) + "\n")); err != nil {
			return InError(err)
		}
		return Done()
	}

	if _, err := p.file.Write([]byte("##" + Serialise(list) + "\n")); err != nil {
		return InError(err)
	}
	return Done()
}

<|endoftext|>


<commit_before>package presence

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"os"
	"path"
	"runtime"
	"strings"

	"github.com/eikeon/marvin/nog"
	"github.com/eikeon/presence"
)

var Root = ""

func init() {
	_, filename, _, _ := runtime.Caller(0)
	Root = path.Dir(filename)
}

type Presence struct {
	nog.InOut
	Present map[string]bool
}

func (s *Presence) Run(in <-chan nog.Message, out chan<- nog.Message) {
	var presenceChannel chan presence.Presence

	name := "presence.html"
	if j, err := os.OpenFile(path.Join(Root, name), os.O_RDONLY, 0666); err == nil {
		if b, err := ioutil.ReadAll(j); err == nil {
			out <- nog.NewMessage("Marvin", string(b), "template")
		} else {
			log.Println("ERROR reading:", err)
		}
	} else {
		log.Println("WARNING: could not open ", name, err)
	}

	for {
		select {
		case m := <-in:
			if m.Why == "statechanged" {
				dec := json.NewDecoder(strings.NewReader(m.What))
				if err := dec.Decode(s); err != nil {
					log.Println("presence decode err:", err)
				}
				if presenceChannel == nil {
					presenceChannel = presence.Listen(s.Present)
				}
			}
		case p := <-presenceChannel:
			if s.Present[p.Name] != p.Status {
				s.Present[p.Name] = p.Status
				var status string
				if p.Status {
					status = "home"
				} else {
					status = "away"
				}
				out <- nog.NewMessage("Marvin", p.Name+" is "+status, "Presence")

				if what, err := json.Marshal(s); err == nil {
					out <- nog.NewMessage("Marvin", string(what), "statechanged")
				} else {
					log.Println("StateChanged err:", err)
				}
			}
		}
	}
}
<commit_msg>Fix for when presence map does not yet exist.<commit_after>package presence

import (
	"encoding/json"
	"io/ioutil"
	"log"
	"os"
	"path"
	"runtime"
	"strings"

	"github.com/eikeon/marvin/nog"
	"github.com/eikeon/presence"
)

var Root = ""

func init() {
	_, filename, _, _ := runtime.Caller(0)
	Root = path.Dir(filename)
}

type Presence struct {
	nog.InOut
	Present map[string]bool
}

func (s *Presence) Run(in <-chan nog.Message, out chan<- nog.Message) {
	var presenceChannel chan presence.Presence

	name := "presence.html"
	if j, err := os.OpenFile(path.Join(Root, name), os.O_RDONLY, 0666); err == nil {
		if b, err := ioutil.ReadAll(j); err == nil {
			out <- nog.NewMessage("Marvin", string(b), "template")
		} else {
			log.Println("ERROR reading:", err)
		}
	} else {
		log.Println("WARNING: could not open ", name, err)
	}

	for {
		select {
		case m := <-in:
			if m.Why == "statechanged" {
				dec := json.NewDecoder(strings.NewReader(m.What))
				if err := dec.Decode(s); err != nil {
					log.Println("presence decode err:", err)
				}
				if s.Present == nil {
					s.Present = make(map[string]bool)
				}
				if presenceChannel == nil {
					presenceChannel = presence.Listen(s.Present)
				}
			}
		case p := <-presenceChannel:
			if s.Present[p.Name] != p.Status {
				s.Present[p.Name] = p.Status
				var status string
				if p.Status {
					status = "home"
				} else {
					status = "away"
				}
				out <- nog.NewMessage("Marvin", p.Name+" is "+status, "Presence")

				if what, err := json.Marshal(s); err == nil {
					out <- nog.NewMessage("Marvin", string(what), "statechanged")
				} else {
					log.Println("StateChanged err:", err)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"path/filepath"
	"time"
)

const (
	NUM_WORKERS_PROD      int = 100
	MASTER_NAME               = "build101-m5"
	WORKER_NAME_TEMPLATE      = "build%d-m5"
	GS_HTTP_LINK              = "https://storage.cloud.google.com/"
	LOGS_LINK_PREFIX          = "http://uberchromegw.corp.google.com/i/skia-ct-worker"
	CT_EMAIL_DISPLAY_NAME     = "Cluster Telemetry"

	// File names and dir names.
	TIMESTAMP_FILE_NAME          = "TIMESTAMP"
	CHROMIUM_BUILDS_DIR_NAME     = "chromium_builds"
	PAGESETS_DIR_NAME            = "page_sets"
	WEB_ARCHIVES_DIR_NAME        = "webpage_archives"
	SKPS_DIR_NAME                = "skps"
	STORAGE_DIR_NAME             = "storage"
	REPO_DIR_NAME                = "skia-repo"
	TASKS_DIR_NAME               = "tasks"
	LUA_TASKS_DIR_NAME           = "lua_runs"
	BENCHMARK_TASKS_DIR_NAME     = "benchmark_runs"
	CHROMIUM_PERF_TASKS_DIR_NAME = "chromium_perf_runs"
	FIX_ARCHIVE_TASKS_DIR_NAME   = "fix_archive_runs"

	// Limit the number of times CT tries to get a remote file before giving up.
	MAX_URI_GET_TRIES = 4

	// Activity constants.
	ACTIVITY_CREATING_PAGESETS        = "CREATING_PAGESETS"
	ACTIVITY_CAPTURING_ARCHIVES       = "CAPTURING_ARCHIVES"
	ACTIVITY_CAPTURING_SKPS           = "CAPTURING_SKPS"
	ACTIVITY_RUNNING_LUA_SCRIPTS      = "RUNNING_LUA_SCRIPTS"
	ACTIVITY_RUNNING_CHROMIUM_PERF    = "RUNNING_CHROMIUM_PERF"
	ACTIVITY_RUNNING_SKIA_CORRECTNESS = "RUNNING_SKIA_CORRECTNESS"
	ACTIVITY_FIXING_ARCHIVES          = "FIXING_ARCHIVES"

	// Pageset types supported by CT.
	PAGESET_TYPE_ALL        = "All"
	PAGESET_TYPE_10k        = "10k"
	PAGESET_TYPE_MOBILE_10k = "Mobile10k"
	PAGESET_TYPE_DUMMY_1k   = "Dummy1k" // Used for testing.

	// Names of binaries executed by CT.
	BINARY_CHROME          = "chrome"
	BINARY_RECORD_WPR      = "record_wpr"
	BINARY_RUN_BENCHMARK   = "run_benchmark"
	BINARY_GCLIENT         = "gclient"
	BINARY_MAKE            = "make"
	BINARY_LUA_PICTURES    = "lua_pictures"
	BINARY_ADB             = "adb"
	BINARY_GIT             = "git"
	BINARY_RENDER_PICTURES = "render_pictures"
	BINARY_MAIL            = "mail"
	BINARY_LUA             = "lua"

	// Platforms supported by CT.
	PLATFORM_ANDROID = "Android"
	PLATFORM_LINUX   = "Linux"

	// Benchmarks supported by CT.
	BENCHMARK_SKPICTURE_PRINTER = "skpicture_printer"
	BENCHMARK_RR                = "rasterize_and_record_micro"
	BENCHMARK_REPAINT           = "repaint"
	BENCHMARK_SMOOTHNESS        = "smoothness"

	// Logserver links. These are only accessible from Google corp.
	MASTER_LOGSERVER_LINK  = "http://uberchromegw.corp.google.com/i/skia-ct-master/"
	WORKERS_LOGSERVER_LINK = "http://uberchromegw.corp.google.com/i/skia-ct-master/all_logs"

	// Default browser args when running benchmarks.
	DEFAULT_BROWSER_ARGS = "--disable-setuid-sandbox --enable-threaded-compositing --enable-impl-side-painting"

	// Timeouts

	PKILL_TIMEOUT = 5 * time.Minute

	// util.SyncDir
	GIT_PULL_TIMEOUT     = 10 * time.Minute
	GCLIENT_SYNC_TIMEOUT = 15 * time.Minute

	// util.BuildSkiaTools
	MAKE_CLEAN_TIMEOUT = 5 * time.Minute
	MAKE_TOOLS_TIMEOUT = 5 * time.Minute

	// util.ResetCheckout
	GIT_RESET_TIMEOUT = 5 * time.Minute
	GIT_CLEAN_TIMEOUT = 5 * time.Minute
	// util.resetChromiumCheckout calls ResetCheckout three times.
	RESET_CHROMIUM_CHECKOUT_TIMEOUT = 3 * (GIT_RESET_TIMEOUT + GIT_CLEAN_TIMEOUT)

	// util.CreateChromiumBuild
	SYNC_SKIA_IN_CHROME_TIMEOUT   = 2 * time.Hour
	GIT_LS_REMOTE_TIMEOUT         = 5 * time.Minute
	GIT_APPLY_TIMEOUT             = 5 * time.Minute
	GOMA_CTL_RESTART_TIMEOUT      = 10 * time.Minute
	GYP_CHROMIUM_TIMEOUT          = 30 * time.Minute
	NINJA_TIMEOUT                 = 2 * time.Hour
	CREATE_CHROMIUM_BUILD_TIMEOUT = SYNC_SKIA_IN_CHROME_TIMEOUT + GIT_LS_REMOTE_TIMEOUT +
		// Three patches are applied when applyPatches is specified.
		3*GIT_APPLY_TIMEOUT +
		// The build steps are repeated twice when applyPatches is specified.
		2*(GOMA_CTL_RESTART_TIMEOUT+GYP_CHROMIUM_TIMEOUT+NINJA_TIMEOUT+
			RESET_CHROMIUM_CHECKOUT_TIMEOUT)

	// util.InstallChromeAPK
	ADB_INSTALL_TIMEOUT = 15 * time.Minute

	// Allow extra time for updating frontend and any other computation not included in the
	// worker timeouts.
	MASTER_SCRIPT_TIMEOUT_PADDING = 30 * time.Minute

	// Build Chromium Task
	GIT_LOG_TIMEOUT                      = 5 * time.Minute
	MASTER_SCRIPT_BUILD_CHROMIUM_TIMEOUT = CREATE_CHROMIUM_BUILD_TIMEOUT + GIT_LOG_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Capture Archives
	// Setting a 5 day timeout since it may take a while to capture 1M archives.
	CAPTURE_ARCHIVES_DEFAULT_CT_BENCHMARK  = "rasterize_and_record_micro_ct"
	CAPTURE_ARCHIVES_TIMEOUT               = 5 * 24 * time.Hour
	MASTER_SCRIPT_CAPTURE_ARCHIVES_TIMEOUT = CAPTURE_ARCHIVES_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Capture SKPs
	REMOVE_INVALID_SKPS_TIMEOUT = 3 * time.Hour
	// Setting a 2 day timeout since it may take a while to capture 1M SKPs.
	CAPTURE_SKPS_TIMEOUT               = 2 * 24 * time.Hour
	MASTER_SCRIPT_CAPTURE_SKPS_TIMEOUT = CAPTURE_SKPS_TIMEOUT + MASTER_SCRIPT_TIMEOUT_PADDING

	// Check Workers Health
	ADB_DEVICES_TIMEOUT          = 30 * time.Minute
	ADB_SHELL_UPTIME_TIMEOUT     = 30 * time.Minute
	CHECK_WORKERS_HEALTH_TIMEOUT = ADB_DEVICES_TIMEOUT + ADB_SHELL_UPTIME_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Create Pagesets
	// Setting a 4 hour timeout since it may take a while to upload page sets to
	// Google Storage when doing 10k page sets per worker.
	CREATE_PAGESETS_TIMEOUT               = 4 * time.Hour
	MASTER_SCRIPT_CREATE_PAGESETS_TIMEOUT = CREATE_PAGESETS_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Run Chromium Perf
	ADB_VERSION_TIMEOUT            = 5 * time.Minute
	ADB_ROOT_TIMEOUT               = 5 * time.Minute
	CSV_PIVOT_TABLE_MERGER_TIMEOUT = 10 * time.Minute
	REBOOT_TIMEOUT                 = 5 * time.Minute
	CSV_MERGER_TIMEOUT             = 1 * time.Hour
	CSV_COMPARER_TIMEOUT           = 2 * time.Hour
	// Setting a 1 day timeout since it may take a while run benchmarks with many
	// repeats.
	RUN_CHROMIUM_PERF_TIMEOUT = 1 * 24 * time.Hour
	// csv_merger runs once for nopatch and once for withpatch
	MASTER_SCRIPT_RUN_CHROMIUM_PERF_TIMEOUT = CREATE_CHROMIUM_BUILD_TIMEOUT + REBOOT_TIMEOUT +
		RUN_CHROMIUM_PERF_TIMEOUT + 2*CSV_MERGER_TIMEOUT + CSV_COMPARER_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Run Lua
	LUA_PICTURES_TIMEOUT          = 2 * time.Hour
	RUN_LUA_TIMEOUT               = 2 * time.Hour
	LUA_AGGREGATOR_TIMEOUT        = 1 * time.Hour
	MASTER_SCRIPT_RUN_LUA_TIMEOUT = RUN_LUA_TIMEOUT + LUA_AGGREGATOR_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Fix Archives
	// Setting a 1 day timeout since it may take a while to validate archives.
	FIX_ARCHIVES_TIMEOUT = 1 * 24 * time.Hour

	// Poller
	MAKE_ALL_TIMEOUT = 15 * time.Minute

	WEBHOOK_SALT_MSG = `For prod, set this file to the value of GCE metadata key webhook_request_salt or call webhook.MustInitRequestSaltFromMetadata() if running in GCE. For testing, run 'echo -n "notverysecret" | base64 -w 0 > /b/storage/webhook_salt.data' or call frontend.InitForTesting().`
)

type PagesetTypeInfo struct {
	NumPages                   int
	CSVSource                  string
	UserAgent                  string
	CaptureArchivesTimeoutSecs int
	CreatePagesetsTimeoutSecs  int
	CaptureSKPsTimeoutSecs     int
	RunChromiumPerfTimeoutSecs int
	Description                string
}

var (
	Master       = fmt.Sprintf(WORKER_NAME_TEMPLATE, 101)
	CtUser       = "chrome-bot"
	Slaves       = GetCTWorkersProd()
	GSBucketName = "cluster-telemetry"

	// Email address of cluster telemetry admins. They will be notified everytime
	// a task has started and completed.
	CtAdmins = []string{"rmistry@google.com"}

	// Names of local directories and files.
	StorageDir           = filepath.Join("/", "b", STORAGE_DIR_NAME)
	RepoDir              = filepath.Join("/", "b", REPO_DIR_NAME)
	GomaDir              = filepath.Join("/", "b", "build", "goma")
	ChromiumBuildsDir    = filepath.Join(StorageDir, CHROMIUM_BUILDS_DIR_NAME)
	ChromiumSrcDir       = filepath.Join(StorageDir, "chromium", "src")
	TelemetryBinariesDir = filepath.Join(ChromiumSrcDir, "tools", "perf")
	TelemetrySrcDir      = filepath.Join(ChromiumSrcDir, "tools", "telemetry")
	CatapultSrcDir       = filepath.Join(ChromiumSrcDir, "third_party", "catapult")
	TaskFileDir          = filepath.Join(StorageDir, "current_task")
	ClientSecretPath     = filepath.Join(StorageDir, "client_secret.json")
	GSTokenPath          = filepath.Join(StorageDir, "google_storage_token.data")
	EmailTokenPath       = filepath.Join(StorageDir, "email.data")
	WebappPasswordPath   = filepath.Join(StorageDir, "webapp.data")
	// Salt used to authenticate webhook requests, base64-encoded. See WEBHOOK_SALT_MSG.
	WebhookRequestSaltPath = filepath.Join(StorageDir, "webhook_salt.data")
	PagesetsDir            = filepath.Join(StorageDir, PAGESETS_DIR_NAME)
	WebArchivesDir         = filepath.Join(StorageDir, WEB_ARCHIVES_DIR_NAME)
	SkpsDir                = filepath.Join(StorageDir, SKPS_DIR_NAME)
	GLogDir                = filepath.Join(StorageDir, "glog")
	ApkName                = "ChromePublic.apk"
	SkiaTreeDir            = filepath.Join(RepoDir, "trunk")
	CtTreeDir              = filepath.Join(RepoDir, "go", "src", "go.skia.org", "infra", "ct")

	// Names of remote directories and files.
	LuaRunsDir          = filepath.Join(TASKS_DIR_NAME, LUA_TASKS_DIR_NAME)
	BenchmarkRunsDir    = filepath.Join(TASKS_DIR_NAME, BENCHMARK_TASKS_DIR_NAME)
	ChromiumPerfRunsDir = filepath.Join(TASKS_DIR_NAME, CHROMIUM_PERF_TASKS_DIR_NAME)
	FixArchivesRunsDir  = filepath.Join(TASKS_DIR_NAME, FIX_ARCHIVE_TASKS_DIR_NAME)

	// Map CT benchmarks to the names recognized by Telemetry.
	BenchmarksToTelemetryName = map[string]string{
		BENCHMARK_SKPICTURE_PRINTER: "skpicture_printer_ct",
		BENCHMARK_RR:                "rasterize_and_record_micro_ct",
		BENCHMARK_REPAINT:           "repaint_ct",
	}

	// Information about the different CT pageset types.
	PagesetTypeToInfo = map[string]*PagesetTypeInfo{
		PAGESET_TYPE_ALL: &PagesetTypeInfo{
			NumPages:                   1000000,
			CSVSource:                  "csv/top-1m.csv",
			UserAgent:                  "desktop",
			CreatePagesetsTimeoutSecs:  60,
			CaptureArchivesTimeoutSecs: 300,
			CaptureSKPsTimeoutSecs:     300,
			RunChromiumPerfTimeoutSecs: 300,
			Description:                "Top 1M (with desktop user-agent)",
		},
		PAGESET_TYPE_10k: &PagesetTypeInfo{
			NumPages:                   10000,
			CSVSource:                  "csv/top-1m.csv",
			UserAgent:                  "desktop",
			CreatePagesetsTimeoutSecs:  60,
			CaptureArchivesTimeoutSecs: 300,
			CaptureSKPsTimeoutSecs:     300,
			RunChromiumPerfTimeoutSecs: 300,
			Description:                "Top 10K (with desktop user-agent)",
		},
		PAGESET_TYPE_MOBILE_10k: &PagesetTypeInfo{
			NumPages:                   10000,
			CSVSource:                  "csv/android-top-1m.csv",
			UserAgent:                  "mobile",
			CreatePagesetsTimeoutSecs:  60,
			CaptureArchivesTimeoutSecs: 300,
			CaptureSKPsTimeoutSecs:     300,
			RunChromiumPerfTimeoutSecs: 300,
			Description:                "Top 10K (with mobile user-agent)",
		},
		PAGESET_TYPE_DUMMY_1k: &PagesetTypeInfo{
			NumPages:                   1000,
			CSVSource:                  "csv/android-top-1m.csv",
			UserAgent:                  "mobile",
			CreatePagesetsTimeoutSecs:  60,
			CaptureArchivesTimeoutSecs: 300,
			CaptureSKPsTimeoutSecs:     300,
			RunChromiumPerfTimeoutSecs: 300,
			Description:                "Top 1K (used for testing, hidden from Runs History by default)",
		},
	}

	// Frontend constants below.
	SupportedBenchmarks = []string{
		BENCHMARK_RR,
		BENCHMARK_REPAINT,
	}

	SupportedPlatformsToDesc = map[string]string{
		PLATFORM_LINUX:   "Linux (100 Ubuntu12.04 machines)",
		PLATFORM_ANDROID: "Android (100 N5 devices)",
	}

	SupportedPageSetsToDesc = map[string]string{
		PLATFORM_LINUX:   "Linux (100 Ubuntu12.04 machines)",
		PLATFORM_ANDROID: "Android (100 N5 devices)",
	}
)

func NumWorkers() int {
	return len(Slaves)
}
<commit_msg>Add benjaminwagner to CtAdmins to watch canaries.<commit_after>package util

import (
	"fmt"
	"path/filepath"
	"time"
)

const (
	NUM_WORKERS_PROD      int = 100
	MASTER_NAME               = "build101-m5"
	WORKER_NAME_TEMPLATE      = "build%d-m5"
	GS_HTTP_LINK              = "https://storage.cloud.google.com/"
	LOGS_LINK_PREFIX          = "http://uberchromegw.corp.google.com/i/skia-ct-worker"
	CT_EMAIL_DISPLAY_NAME     = "Cluster Telemetry"

	// File names and dir names.
	TIMESTAMP_FILE_NAME          = "TIMESTAMP"
	CHROMIUM_BUILDS_DIR_NAME     = "chromium_builds"
	PAGESETS_DIR_NAME            = "page_sets"
	WEB_ARCHIVES_DIR_NAME        = "webpage_archives"
	SKPS_DIR_NAME                = "skps"
	STORAGE_DIR_NAME             = "storage"
	REPO_DIR_NAME                = "skia-repo"
	TASKS_DIR_NAME               = "tasks"
	LUA_TASKS_DIR_NAME           = "lua_runs"
	BENCHMARK_TASKS_DIR_NAME     = "benchmark_runs"
	CHROMIUM_PERF_TASKS_DIR_NAME = "chromium_perf_runs"
	FIX_ARCHIVE_TASKS_DIR_NAME   = "fix_archive_runs"

	// Limit the number of times CT tries to get a remote file before giving up.
	MAX_URI_GET_TRIES = 4

	// Activity constants.
	ACTIVITY_CREATING_PAGESETS        = "CREATING_PAGESETS"
	ACTIVITY_CAPTURING_ARCHIVES       = "CAPTURING_ARCHIVES"
	ACTIVITY_CAPTURING_SKPS           = "CAPTURING_SKPS"
	ACTIVITY_RUNNING_LUA_SCRIPTS      = "RUNNING_LUA_SCRIPTS"
	ACTIVITY_RUNNING_CHROMIUM_PERF    = "RUNNING_CHROMIUM_PERF"
	ACTIVITY_RUNNING_SKIA_CORRECTNESS = "RUNNING_SKIA_CORRECTNESS"
	ACTIVITY_FIXING_ARCHIVES          = "FIXING_ARCHIVES"

	// Pageset types supported by CT.
	PAGESET_TYPE_ALL        = "All"
	PAGESET_TYPE_10k        = "10k"
	PAGESET_TYPE_MOBILE_10k = "Mobile10k"
	PAGESET_TYPE_DUMMY_1k   = "Dummy1k" // Used for testing.

	// Names of binaries executed by CT.
	BINARY_CHROME          = "chrome"
	BINARY_RECORD_WPR      = "record_wpr"
	BINARY_RUN_BENCHMARK   = "run_benchmark"
	BINARY_GCLIENT         = "gclient"
	BINARY_MAKE            = "make"
	BINARY_LUA_PICTURES    = "lua_pictures"
	BINARY_ADB             = "adb"
	BINARY_GIT             = "git"
	BINARY_RENDER_PICTURES = "render_pictures"
	BINARY_MAIL            = "mail"
	BINARY_LUA             = "lua"

	// Platforms supported by CT.
	PLATFORM_ANDROID = "Android"
	PLATFORM_LINUX   = "Linux"

	// Benchmarks supported by CT.
	BENCHMARK_SKPICTURE_PRINTER = "skpicture_printer"
	BENCHMARK_RR                = "rasterize_and_record_micro"
	BENCHMARK_REPAINT           = "repaint"
	BENCHMARK_SMOOTHNESS        = "smoothness"

	// Logserver links. These are only accessible from Google corp.
	MASTER_LOGSERVER_LINK  = "http://uberchromegw.corp.google.com/i/skia-ct-master/"
	WORKERS_LOGSERVER_LINK = "http://uberchromegw.corp.google.com/i/skia-ct-master/all_logs"

	// Default browser args when running benchmarks.
	DEFAULT_BROWSER_ARGS = "--disable-setuid-sandbox --enable-threaded-compositing --enable-impl-side-painting"

	// Timeouts

	PKILL_TIMEOUT = 5 * time.Minute

	// util.SyncDir
	GIT_PULL_TIMEOUT     = 10 * time.Minute
	GCLIENT_SYNC_TIMEOUT = 15 * time.Minute

	// util.BuildSkiaTools
	MAKE_CLEAN_TIMEOUT = 5 * time.Minute
	MAKE_TOOLS_TIMEOUT = 5 * time.Minute

	// util.ResetCheckout
	GIT_RESET_TIMEOUT = 5 * time.Minute
	GIT_CLEAN_TIMEOUT = 5 * time.Minute
	// util.resetChromiumCheckout calls ResetCheckout three times.
	RESET_CHROMIUM_CHECKOUT_TIMEOUT = 3 * (GIT_RESET_TIMEOUT + GIT_CLEAN_TIMEOUT)

	// util.CreateChromiumBuild
	SYNC_SKIA_IN_CHROME_TIMEOUT   = 2 * time.Hour
	GIT_LS_REMOTE_TIMEOUT         = 5 * time.Minute
	GIT_APPLY_TIMEOUT             = 5 * time.Minute
	GOMA_CTL_RESTART_TIMEOUT      = 10 * time.Minute
	GYP_CHROMIUM_TIMEOUT          = 30 * time.Minute
	NINJA_TIMEOUT                 = 2 * time.Hour
	CREATE_CHROMIUM_BUILD_TIMEOUT = SYNC_SKIA_IN_CHROME_TIMEOUT + GIT_LS_REMOTE_TIMEOUT +
		// Three patches are applied when applyPatches is specified.
		3*GIT_APPLY_TIMEOUT +
		// The build steps are repeated twice when applyPatches is specified.
		2*(GOMA_CTL_RESTART_TIMEOUT+GYP_CHROMIUM_TIMEOUT+NINJA_TIMEOUT+
			RESET_CHROMIUM_CHECKOUT_TIMEOUT)

	// util.InstallChromeAPK
	ADB_INSTALL_TIMEOUT = 15 * time.Minute

	// Allow extra time for updating frontend and any other computation not included in the
	// worker timeouts.
	MASTER_SCRIPT_TIMEOUT_PADDING = 30 * time.Minute

	// Build Chromium Task
	GIT_LOG_TIMEOUT                      = 5 * time.Minute
	MASTER_SCRIPT_BUILD_CHROMIUM_TIMEOUT = CREATE_CHROMIUM_BUILD_TIMEOUT + GIT_LOG_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Capture Archives
	// Setting a 5 day timeout since it may take a while to capture 1M archives.
	CAPTURE_ARCHIVES_DEFAULT_CT_BENCHMARK  = "rasterize_and_record_micro_ct"
	CAPTURE_ARCHIVES_TIMEOUT               = 5 * 24 * time.Hour
	MASTER_SCRIPT_CAPTURE_ARCHIVES_TIMEOUT = CAPTURE_ARCHIVES_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Capture SKPs
	REMOVE_INVALID_SKPS_TIMEOUT = 3 * time.Hour
	// Setting a 2 day timeout since it may take a while to capture 1M SKPs.
	CAPTURE_SKPS_TIMEOUT               = 2 * 24 * time.Hour
	MASTER_SCRIPT_CAPTURE_SKPS_TIMEOUT = CAPTURE_SKPS_TIMEOUT + MASTER_SCRIPT_TIMEOUT_PADDING

	// Check Workers Health
	ADB_DEVICES_TIMEOUT          = 30 * time.Minute
	ADB_SHELL_UPTIME_TIMEOUT     = 30 * time.Minute
	CHECK_WORKERS_HEALTH_TIMEOUT = ADB_DEVICES_TIMEOUT + ADB_SHELL_UPTIME_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Create Pagesets
	// Setting a 4 hour timeout since it may take a while to upload page sets to
	// Google Storage when doing 10k page sets per worker.
	CREATE_PAGESETS_TIMEOUT               = 4 * time.Hour
	MASTER_SCRIPT_CREATE_PAGESETS_TIMEOUT = CREATE_PAGESETS_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Run Chromium Perf
	ADB_VERSION_TIMEOUT            = 5 * time.Minute
	ADB_ROOT_TIMEOUT               = 5 * time.Minute
	CSV_PIVOT_TABLE_MERGER_TIMEOUT = 10 * time.Minute
	REBOOT_TIMEOUT                 = 5 * time.Minute
	CSV_MERGER_TIMEOUT             = 1 * time.Hour
	CSV_COMPARER_TIMEOUT           = 2 * time.Hour
	// Setting a 1 day timeout since it may take a while run benchmarks with many
	// repeats.
	RUN_CHROMIUM_PERF_TIMEOUT = 1 * 24 * time.Hour
	// csv_merger runs once for nopatch and once for withpatch
	MASTER_SCRIPT_RUN_CHROMIUM_PERF_TIMEOUT = CREATE_CHROMIUM_BUILD_TIMEOUT + REBOOT_TIMEOUT +
		RUN_CHROMIUM_PERF_TIMEOUT + 2*CSV_MERGER_TIMEOUT + CSV_COMPARER_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Run Lua
	LUA_PICTURES_TIMEOUT          = 2 * time.Hour
	RUN_LUA_TIMEOUT               = 2 * time.Hour
	LUA_AGGREGATOR_TIMEOUT        = 1 * time.Hour
	MASTER_SCRIPT_RUN_LUA_TIMEOUT = RUN_LUA_TIMEOUT + LUA_AGGREGATOR_TIMEOUT +
		MASTER_SCRIPT_TIMEOUT_PADDING

	// Fix Archives
	// Setting a 1 day timeout since it may take a while to validate archives.
	FIX_ARCHIVES_TIMEOUT = 1 * 24 * time.Hour

	// Poller
	MAKE_ALL_TIMEOUT = 15 * time.Minute

	WEBHOOK_SALT_MSG = `For prod, set this file to the value of GCE metadata key webhook_request_salt or call webhook.MustInitRequestSaltFromMetadata() if running in GCE. For testing, run 'echo -n "notverysecret" | base64 -w 0 > /b/storage/webhook_salt.data' or call frontend.InitForTesting().`
)

type PagesetTypeInfo struct {
	NumPages                   int
	CSVSource                  string
	UserAgent                  string
	CaptureArchivesTimeoutSecs int
	CreatePagesetsTimeoutSecs  int
	CaptureSKPsTimeoutSecs     int
	RunChromiumPerfTimeoutSecs int
	Description                string
}

var (
	Master       = fmt.Sprintf(WORKER_NAME_TEMPLATE, 101)
	CtUser       = "chrome-bot"
	Slaves       = GetCTWorkersProd()
	GSBucketName = "cluster-telemetry"

	// Email address of cluster telemetry admins. They will be notified everytime
	// a task has started and completed.
	CtAdmins = []string{"rmistry@google.com", "benjaminwagner@google.com"}

	// Names of local directories and files.
	StorageDir           = filepath.Join("/", "b", STORAGE_DIR_NAME)
	RepoDir              = filepath.Join("/", "b", REPO_DIR_NAME)
	GomaDir              = filepath.Join("/", "b", "build", "goma")
	ChromiumBuildsDir    = filepath.Join(StorageDir, CHROMIUM_BUILDS_DIR_NAME)
	ChromiumSrcDir       = filepath.Join(StorageDir, "chromium", "src")
	TelemetryBinariesDir = filepath.Join(ChromiumSrcDir, "tools", "perf")
	TelemetrySrcDir      = filepath.Join(ChromiumSrcDir, "tools", "telemetry")
	CatapultSrcDir       = filepath.Join(ChromiumSrcDir, "third_party", "catapult")
	TaskFileDir          = filepath.Join(StorageDir, "current_task")
	ClientSecretPath     = filepath.Join(StorageDir, "client_secret.json")
	GSTokenPath          = filepath.Join(StorageDir, "google_storage_token.data")
	EmailTokenPath       = filepath.Join(StorageDir, "email.data")
	WebappPasswordPath   = filepath.Join(StorageDir, "webapp.data")
	// Salt used to authenticate webhook requests, base64-encoded. See WEBHOOK_SALT_MSG.
	WebhookRequestSaltPath = filepath.Join(StorageDir, "webhook_salt.data")
	PagesetsDir            = filepath.Join(StorageDir, PAGESETS_DIR_NAME)
	WebArchivesDir         = filepath.Join(StorageDir, WEB_ARCHIVES_DIR_NAME)
	SkpsDir                = filepath.Join(StorageDir, SKPS_DIR_NAME)
	GLogDir                = filepath.Join(StorageDir, "glog")
	ApkName                = "ChromePublic.apk"
	SkiaTreeDir            = filepath.Join(RepoDir, "trunk")
	CtTreeDir              = filepath.Join(RepoDir, "go", "src", "go.skia.org", "infra", "ct")

	// Names of remote directories and files.
	LuaRunsDir          = filepath.Join(TASKS_DIR_NAME, LUA_TASKS_DIR_NAME)
	BenchmarkRunsDir    = filepath.Join(TASKS_DIR_NAME, BENCHMARK_TASKS_DIR_NAME)
	ChromiumPerfRunsDir = filepath.Join(TASKS_DIR_NAME, CHROMIUM_PERF_TASKS_DIR_NAME)
	FixArchivesRunsDir  = filepath.Join(TASKS_DIR_NAME, FIX_ARCHIVE_TASKS_DIR_NAME)

	// Map CT benchmarks to the names recognized by Telemetry.
	BenchmarksToTelemetryName = map[string]string{
		BENCHMARK_SKPICTURE_PRINTER: "skpicture_printer_ct",
		BENCHMARK_RR:                "rasterize_and_record_micro_ct",
		BENCHMARK_REPAINT:           "repaint_ct",
	}

	// Information about the different CT pageset types.
	PagesetTypeToInfo = map[string]*PagesetTypeInfo{
		PAGESET_TYPE_ALL: &PagesetTypeInfo{
			NumPages:                   1000000,
			CSVSource:                  "csv/top-1m.csv",
			UserAgent:                  "desktop",
			CreatePagesetsTimeoutSecs:  60,
			CaptureArchivesTimeoutSecs: 300,
			CaptureSKPsTimeoutSecs:     300,
			RunChromiumPerfTimeoutSecs: 300,
			Description:                "Top 1M (with desktop user-agent)",
		},
		PAGESET_TYPE_10k: &PagesetTypeInfo{
			NumPages:                   10000,
			CSVSource:                  "csv/top-1m.csv",
			UserAgent:                  "desktop",
			CreatePagesetsTimeoutSecs:  60,
			CaptureArchivesTimeoutSecs: 300,
			CaptureSKPsTimeoutSecs:     300,
			RunChromiumPerfTimeoutSecs: 300,
			Description:                "Top 10K (with desktop user-agent)",
		},
		PAGESET_TYPE_MOBILE_10k: &PagesetTypeInfo{
			NumPages:                   10000,
			CSVSource:                  "csv/android-top-1m.csv",
			UserAgent:                  "mobile",
			CreatePagesetsTimeoutSecs:  60,
			CaptureArchivesTimeoutSecs: 300,
			CaptureSKPsTimeoutSecs:     300,
			RunChromiumPerfTimeoutSecs: 300,
			Description:                "Top 10K (with mobile user-agent)",
		},
		PAGESET_TYPE_DUMMY_1k: &PagesetTypeInfo{
			NumPages:                   1000,
			CSVSource:                  "csv/android-top-1m.csv",
			UserAgent:                  "mobile",
			CreatePagesetsTimeoutSecs:  60,
			CaptureArchivesTimeoutSecs: 300,
			CaptureSKPsTimeoutSecs:     300,
			RunChromiumPerfTimeoutSecs: 300,
			Description:                "Top 1K (used for testing, hidden from Runs History by default)",
		},
	}

	// Frontend constants below.
	SupportedBenchmarks = []string{
		BENCHMARK_RR,
		BENCHMARK_REPAINT,
	}

	SupportedPlatformsToDesc = map[string]string{
		PLATFORM_LINUX:   "Linux (100 Ubuntu12.04 machines)",
		PLATFORM_ANDROID: "Android (100 N5 devices)",
	}

	SupportedPageSetsToDesc = map[string]string{
		PLATFORM_LINUX:   "Linux (100 Ubuntu12.04 machines)",
		PLATFORM_ANDROID: "Android (100 N5 devices)",
	}
)

func NumWorkers() int {
	return len(Slaves)
}
<|endoftext|>


<commit_before>package ploop

import (
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

// Ploop is a class representing a ploop
type Ploop struct {
	dd string
}

// Possible SetVerboseLevel arguments
const (
	unsetVerbosity int = -111
	NoConsole          = -2
	NoStdout           = -1
	Timestamps         = 4
	ShowCommands       = 5
)

var verbosity = unsetVerbosity
var verbosityOpt = []string{""}

// SetVerboseLevel sets a level of verbosity when logging to stdout/stderr
func SetVerboseLevel(v int) {
	verbosity = v
	verbosityOpt = []string{"-v" + strconv.Itoa(verbosity)}
}

var once sync.Once

// load ploop modules
func loadKmod() {
	// try to load ploop modules
	modules := []string{"ploop", "pfmt_ploop1", "pfmt_raw", "pio_direct", "pio_nfs", "pio_kaio"}
	for _, m := range modules {
		exec.Command("modprobe", m).Run()
	}
}

// Open opens a ploop DiskDescriptor.xml, most ploop operations require it
func Open(file string) (Ploop, error) {
	var d Ploop

	once.Do(loadKmod)

	d.dd = file
	return d, nil
}

// Close closes a ploop disk descriptor when it is no longer needed
func (d Ploop) Close() {
	d.dd = ""
}

// ImageMode is a type for CreateParam.Mode field
type ImageMode string

// Possible values for ImageMode
const (
	Expanded     ImageMode = "expanded"
	Preallocated ImageMode = "preallocated"
	Raw          ImageMode = "raw"
)

// ParseImageMode converts a string to ImageMode value
func ParseImageMode(s string) (ImageMode, error) {
	switch strings.ToLower(s) {
	case "expanded":
		return Expanded, nil
	case "preallocated":
		return Preallocated, nil
	case "raw":
		return Raw, nil
	default:
		return Expanded, &Err{c: E_PARAM, s: "ParseImageMode: unknown mode " + s}
	}
}

// String converts an ImageMode value to string
func (m ImageMode) String() string {
	return string(m)
}

// CreateFlags is a type for CreateParam.Flags
type CreateFlags int

// Possible values for CreateFlags
const (
	NoLazy CreateFlags = 1 << iota
)

// CreateParam is a set of parameters for a newly created ploop
type CreateParam struct {
	Size  uint64      // image size, in kilobytes (FS size is about 10% smaller)
	Mode  ImageMode   // image mode
	File  string      // path to and a file name for base delta image
	CLog  uint        // cluster block size log (6 to 15, default 11)
	Flags CreateFlags // flags
}

// Create creates a ploop image and its DiskDescriptor.xml
func Create(p *CreateParam) error {
	once.Do(loadKmod)

	// default image file name
	if p.File == "" {
		p.File = "root.hdd"
	}

	args := []string{"init", "-s", strconv.FormatUint(p.Size, 10) + "K"}
	if p.Mode != "" {
		args = append(args, "-m", string(p.Mode))
	}
	if p.CLog != 0 {
		// ploop cluster block size, in 512-byte sectors
		// default is 1M cluster block size (CLog=11)
		// 2^11 = 2048 sectors, 2048*512 = 1M
		blocksize := 1 << p.CLog
		args = append(args, "-b", strconv.Itoa(blocksize))
	}
	if p.Flags != 0 {
		if p.Flags&NoLazy == NoLazy {
			args = append(args, "--nolazy")
		}
	}
	args = append(args, p.File)

	return ploop(args...)
}

// MountParam is a set of parameters to pass to Mount()
type MountParam struct {
	UUID     string // snapshot uuid (empty for top delta)
	Target   string // mount point (empty if no mount is needed)
	Data     string // auxiliary mount options
	Readonly bool   // mount read-only
	Fsck     bool   // do fsck before mounting inner FS
}

// Input is like this (with or without a timestamp, depending on verbosity):
// [ 0.001234] Adding delta dev=/dev/ploop12345 ...
// Adding delta dev=/dev/ploop12345 ...
var reAddDelta = regexp.MustCompile(`(?m)^(?:\[[0-9. ]+\] )*Adding delta dev=(/dev/ploop\d+) `)

// Mount creates a ploop device and (optionally) mounts it
func (d Ploop) Mount(p *MountParam) (string, error) {
	args := []string{"mount"}
	if p.Readonly {
		args = append(args, "-r")
	}
	if p.Fsck {
		args = append(args, "-F")
	}
	if p.Target != "" {
		args = append(args, "-m", p.Target)
	}
	if p.Data != "" {
		args = append(args, "-o", p.Data)
	}
	args = append(args, d.dd)

	dev := ""
	out, err := ploopOut(args...)
	if err == nil {
		// Figure out what device we have
		m := reAddDelta.FindStringSubmatch(out)
		if len(m) != 2 {
			return "", &Err{c: -1,
				s: "Can't parse ploop mount output:\n" + out}
		}
		dev = m[1]
	}

	return dev, err
}

// Umount unmounts the ploop filesystem and dismantles the device
func (d Ploop) Umount() error {
	return ploop("umount", d.dd)
}

// UmountByDevice unmounts the ploop filesystem and dismantles the device.
// Unlike Umount(), this is a lower-level function meaning it can be less
// safe and should generally not be used.
func UmountByDevice(dev string) error {
	return ploop("umount", "-d", dev)
}

// Resize changes the ploop size. Offline flag is ignored
// by this implementation, as ploop tool automatically chooses
// whether do to offline resize (if device is not mounted).
func (d Ploop) Resize(size uint64, offline bool) error {
	return ploop("resize", "-s", strconv.FormatUint(size, 10)+"K", d.dd)
}

// Snapshot creates a ploop snapshot, returning its uuid
func (d Ploop) Snapshot() (string, error) {
	uuid, err := UUID()
	if err != nil {
		return "", err
	}

	return uuid, ploop("snapshot", "-u", uuid, d.dd)
}

// SwitchSnapshot switches to a specified snapshot,
// creates a new empty delta on top of it, and makes it a top one
// (i.e. the one new data will be written to).
// Old top delta (i.e. data modified since the last snapshot) is lost.
func (d Ploop) SwitchSnapshot(uuid string) error {
	return ploop("snapshot-switch", "-u", uuid, d.dd)
}

// DeleteSnapshot deletes a snapshot (merging it down if necessary)
func (d Ploop) DeleteSnapshot(uuid string) error {
	return ploop("snapshot-delete", "-u", uuid, d.dd)
}

// ReplaceFlag is a type for ReplaceParam.Flags field
type ReplaceFlag int

// Possible values for ReplaceParam.Flags field
const (
	_ = iota
	// KeepName renames the new file to old file name after replace;
	// note that if this option is used the old file is removed.
	KeepName ReplaceFlag = iota
)

// ReplaceParam is a set of parameters to Replace()
type ReplaceParam struct {
	File string // new image file name
	// Image to be replaced is specified by either
	// uuid, current file name, or level,
	// in the above order of preference.
	UUID    string
	CurFile string
	Level   int
	Flags   ReplaceFlag
}

// Replace replaces a ploop image to a different (but identical) one
func (d Ploop) Replace(p *ReplaceParam) error {
	args := []string{"replace"}
	if p.UUID != "" {
		args = append(args, "-u", p.UUID)
	} else if p.CurFile != "" {
		args = append(args, "-o", p.CurFile)
	} else {
		args = append(args, "-l", strconv.Itoa(p.Level))
	}

	if p.Flags != 0 {
		if p.Flags&KeepName == KeepName {
			args = append(args, "-k")
		}
	}
	args = append(args, "-i", p.File)
	args = append(args, d.dd)

	return ploop(args...)
}

// device:	/dev/ploop25579
var reDevice = regexp.MustCompile(`(?m)^device:\s+(/dev/ploop\d+)$`)

func (d Ploop) getDevice() (string, error) {
	dev := ""
	out, err := ploopOut("-v-1", "info", "-d", d.dd)
	if err == nil {
		// Figure out what device we have
		m := reDevice.FindStringSubmatch(out)
		if len(m) > 1 {
			dev = m[1]
		}
	}
	return dev, err
}

// IsMounted returns true if ploop is mounted
func (d Ploop) IsMounted() (bool, error) {
	dev, err := d.getDevice()
	return dev != "", err
}

// FSInfoData holds information about ploop inner file system
type FSInfoData struct {
	BlockSize  uint64
	Blocks     uint64
	BlocksFree uint64
	Inodes     uint64
	InodesFree uint64
}

//   resource           Size           Used
//  1k-blocks       10188052          36888
//     inodes         655360             12
var reFSInfo = regexp.MustCompile(`
\s+1k-blocks\s+(\d+)\s+(\d+)
\s+inodes\s+(\d+)\s+(\d+)
`)

// FSInfo gets info of ploop's inner file system
func FSInfo(file string) (FSInfoData, error) {
	once.Do(loadKmod)
	var info FSInfoData

	out, err := ploopOut("-v-1", "info", file)
	if err == nil {
		info.BlockSize = 1024 // ploop info reports in 1-k blocks
		i := reFSInfo.FindStringSubmatch(out)
		if len(i) != 5 {
			return info, &Err{c: -1,
				s: "Can't parse ploop info output:\n" + out}
		}

		info.Blocks, _ = strconv.ParseUint(i[1], 10, 64)
		blocksUsed, _ := strconv.ParseUint(i[2], 10, 64)
		info.BlocksFree = info.Blocks - blocksUsed

		info.Inodes, _ = strconv.ParseUint(i[3], 10, 64)
		inodesUsed, _ := strconv.ParseUint(i[4], 10, 64)
		info.InodesFree = info.Inodes - inodesUsed
	}

	return info, err
}

// ImageInfoData holds information about ploop image
type ImageInfoData struct {
	Blocks    uint64
	BlockSize uint32
	Version   int
}

// size:	20971520
// blocksize:	2048
// fmt_version:	2
var reImageInfo = regexp.MustCompile(`\n*size:\s+(\d+)
blocksize:\s+(\d+)
fmt_version:\s+(\d*)
`)

// ImageInfo gets information about a ploop image
func (d Ploop) ImageInfo() (ImageInfoData, error) {
	var info ImageInfoData

	out, err := ploopOut("-v-1", "info", "-s", d.dd)
	if err == nil {
		i := reImageInfo.FindStringSubmatch(out)
		if len(i) != 4 {
			return info, &Err{c: -1,
				s: "Can't parse ploop info -s output:\n" + out}
		}
		info.Blocks, _ = strconv.ParseUint(i[1], 10, 64)
		bs, _ := strconv.ParseUint(i[2], 10, 32)
		info.BlockSize = uint32(bs)
		info.Version, _ = strconv.Atoi(i[3])
	}

	return info, err
}
<commit_msg>Create(): fix Mode flag<commit_after>package ploop

import (
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

// Ploop is a class representing a ploop
type Ploop struct {
	dd string
}

// Possible SetVerboseLevel arguments
const (
	unsetVerbosity int = -111
	NoConsole          = -2
	NoStdout           = -1
	Timestamps         = 4
	ShowCommands       = 5
)

var verbosity = unsetVerbosity
var verbosityOpt = []string{""}

// SetVerboseLevel sets a level of verbosity when logging to stdout/stderr
func SetVerboseLevel(v int) {
	verbosity = v
	verbosityOpt = []string{"-v" + strconv.Itoa(verbosity)}
}

var once sync.Once

// load ploop modules
func loadKmod() {
	// try to load ploop modules
	modules := []string{"ploop", "pfmt_ploop1", "pfmt_raw", "pio_direct", "pio_nfs", "pio_kaio"}
	for _, m := range modules {
		exec.Command("modprobe", m).Run()
	}
}

// Open opens a ploop DiskDescriptor.xml, most ploop operations require it
func Open(file string) (Ploop, error) {
	var d Ploop

	once.Do(loadKmod)

	d.dd = file
	return d, nil
}

// Close closes a ploop disk descriptor when it is no longer needed
func (d Ploop) Close() {
	d.dd = ""
}

// ImageMode is a type for CreateParam.Mode field
type ImageMode string

// Possible values for ImageMode
const (
	Expanded     ImageMode = "expanded"
	Preallocated ImageMode = "preallocated"
	Raw          ImageMode = "raw"
)

// ParseImageMode converts a string to ImageMode value
func ParseImageMode(s string) (ImageMode, error) {
	switch strings.ToLower(s) {
	case "expanded":
		return Expanded, nil
	case "preallocated":
		return Preallocated, nil
	case "raw":
		return Raw, nil
	default:
		return Expanded, &Err{c: E_PARAM, s: "ParseImageMode: unknown mode " + s}
	}
}

// String converts an ImageMode value to string
func (m ImageMode) String() string {
	return string(m)
}

// CreateFlags is a type for CreateParam.Flags
type CreateFlags int

// Possible values for CreateFlags
const (
	NoLazy CreateFlags = 1 << iota
)

// CreateParam is a set of parameters for a newly created ploop
type CreateParam struct {
	Size  uint64      // image size, in kilobytes (FS size is about 10% smaller)
	Mode  ImageMode   // image mode
	File  string      // path to and a file name for base delta image
	CLog  uint        // cluster block size log (6 to 15, default 11)
	Flags CreateFlags // flags
}

// Create creates a ploop image and its DiskDescriptor.xml
func Create(p *CreateParam) error {
	once.Do(loadKmod)

	// default image file name
	if p.File == "" {
		p.File = "root.hdd"
	}

	args := []string{"init", "-s", strconv.FormatUint(p.Size, 10) + "K"}
	if p.Mode != "" {
		args = append(args, "-f", string(p.Mode))
	}
	if p.CLog != 0 {
		// ploop cluster block size, in 512-byte sectors
		// default is 1M cluster block size (CLog=11)
		// 2^11 = 2048 sectors, 2048*512 = 1M
		blocksize := 1 << p.CLog
		args = append(args, "-b", strconv.Itoa(blocksize))
	}
	if p.Flags != 0 {
		if p.Flags&NoLazy == NoLazy {
			args = append(args, "--nolazy")
		}
	}
	args = append(args, p.File)

	return ploop(args...)
}

// MountParam is a set of parameters to pass to Mount()
type MountParam struct {
	UUID     string // snapshot uuid (empty for top delta)
	Target   string // mount point (empty if no mount is needed)
	Data     string // auxiliary mount options
	Readonly bool   // mount read-only
	Fsck     bool   // do fsck before mounting inner FS
}

// Input is like this (with or without a timestamp, depending on verbosity):
// [ 0.001234] Adding delta dev=/dev/ploop12345 ...
// Adding delta dev=/dev/ploop12345 ...
var reAddDelta = regexp.MustCompile(`(?m)^(?:\[[0-9. ]+\] )*Adding delta dev=(/dev/ploop\d+) `)

// Mount creates a ploop device and (optionally) mounts it
func (d Ploop) Mount(p *MountParam) (string, error) {
	args := []string{"mount"}
	if p.Readonly {
		args = append(args, "-r")
	}
	if p.Fsck {
		args = append(args, "-F")
	}
	if p.Target != "" {
		args = append(args, "-m", p.Target)
	}
	if p.Data != "" {
		args = append(args, "-o", p.Data)
	}
	args = append(args, d.dd)

	dev := ""
	out, err := ploopOut(args...)
	if err == nil {
		// Figure out what device we have
		m := reAddDelta.FindStringSubmatch(out)
		if len(m) != 2 {
			return "", &Err{c: -1,
				s: "Can't parse ploop mount output:\n" + out}
		}
		dev = m[1]
	}

	return dev, err
}

// Umount unmounts the ploop filesystem and dismantles the device
func (d Ploop) Umount() error {
	return ploop("umount", d.dd)
}

// UmountByDevice unmounts the ploop filesystem and dismantles the device.
// Unlike Umount(), this is a lower-level function meaning it can be less
// safe and should generally not be used.
func UmountByDevice(dev string) error {
	return ploop("umount", "-d", dev)
}

// Resize changes the ploop size. Offline flag is ignored
// by this implementation, as ploop tool automatically chooses
// whether do to offline resize (if device is not mounted).
func (d Ploop) Resize(size uint64, offline bool) error {
	return ploop("resize", "-s", strconv.FormatUint(size, 10)+"K", d.dd)
}

// Snapshot creates a ploop snapshot, returning its uuid
func (d Ploop) Snapshot() (string, error) {
	uuid, err := UUID()
	if err != nil {
		return "", err
	}

	return uuid, ploop("snapshot", "-u", uuid, d.dd)
}

// SwitchSnapshot switches to a specified snapshot,
// creates a new empty delta on top of it, and makes it a top one
// (i.e. the one new data will be written to).
// Old top delta (i.e. data modified since the last snapshot) is lost.
func (d Ploop) SwitchSnapshot(uuid string) error {
	return ploop("snapshot-switch", "-u", uuid, d.dd)
}

// DeleteSnapshot deletes a snapshot (merging it down if necessary)
func (d Ploop) DeleteSnapshot(uuid string) error {
	return ploop("snapshot-delete", "-u", uuid, d.dd)
}

// ReplaceFlag is a type for ReplaceParam.Flags field
type ReplaceFlag int

// Possible values for ReplaceParam.Flags field
const (
	_ = iota
	// KeepName renames the new file to old file name after replace;
	// note that if this option is used the old file is removed.
	KeepName ReplaceFlag = iota
)

// ReplaceParam is a set of parameters to Replace()
type ReplaceParam struct {
	File string // new image file name
	// Image to be replaced is specified by either
	// uuid, current file name, or level,
	// in the above order of preference.
	UUID    string
	CurFile string
	Level   int
	Flags   ReplaceFlag
}

// Replace replaces a ploop image to a different (but identical) one
func (d Ploop) Replace(p *ReplaceParam) error {
	args := []string{"replace"}
	if p.UUID != "" {
		args = append(args, "-u", p.UUID)
	} else if p.CurFile != "" {
		args = append(args, "-o", p.CurFile)
	} else {
		args = append(args, "-l", strconv.Itoa(p.Level))
	}

	if p.Flags != 0 {
		if p.Flags&KeepName == KeepName {
			args = append(args, "-k")
		}
	}
	args = append(args, "-i", p.File)
	args = append(args, d.dd)

	return ploop(args...)
}

// device:	/dev/ploop25579
var reDevice = regexp.MustCompile(`(?m)^device:\s+(/dev/ploop\d+)$`)

func (d Ploop) getDevice() (string, error) {
	dev := ""
	out, err := ploopOut("-v-1", "info", "-d", d.dd)
	if err == nil {
		// Figure out what device we have
		m := reDevice.FindStringSubmatch(out)
		if len(m) > 1 {
			dev = m[1]
		}
	}
	return dev, err
}

// IsMounted returns true if ploop is mounted
func (d Ploop) IsMounted() (bool, error) {
	dev, err := d.getDevice()
	return dev != "", err
}

// FSInfoData holds information about ploop inner file system
type FSInfoData struct {
	BlockSize  uint64
	Blocks     uint64
	BlocksFree uint64
	Inodes     uint64
	InodesFree uint64
}

//   resource           Size           Used
//  1k-blocks       10188052          36888
//     inodes         655360             12
var reFSInfo = regexp.MustCompile(`
\s+1k-blocks\s+(\d+)\s+(\d+)
\s+inodes\s+(\d+)\s+(\d+)
`)

// FSInfo gets info of ploop's inner file system
func FSInfo(file string) (FSInfoData, error) {
	once.Do(loadKmod)
	var info FSInfoData

	out, err := ploopOut("-v-1", "info", file)
	if err == nil {
		info.BlockSize = 1024 // ploop info reports in 1-k blocks
		i := reFSInfo.FindStringSubmatch(out)
		if len(i) != 5 {
			return info, &Err{c: -1,
				s: "Can't parse ploop info output:\n" + out}
		}

		info.Blocks, _ = strconv.ParseUint(i[1], 10, 64)
		blocksUsed, _ := strconv.ParseUint(i[2], 10, 64)
		info.BlocksFree = info.Blocks - blocksUsed

		info.Inodes, _ = strconv.ParseUint(i[3], 10, 64)
		inodesUsed, _ := strconv.ParseUint(i[4], 10, 64)
		info.InodesFree = info.Inodes - inodesUsed
	}

	return info, err
}

// ImageInfoData holds information about ploop image
type ImageInfoData struct {
	Blocks    uint64
	BlockSize uint32
	Version   int
}

// size:	20971520
// blocksize:	2048
// fmt_version:	2
var reImageInfo = regexp.MustCompile(`\n*size:\s+(\d+)
blocksize:\s+(\d+)
fmt_version:\s+(\d*)
`)

// ImageInfo gets information about a ploop image
func (d Ploop) ImageInfo() (ImageInfoData, error) {
	var info ImageInfoData

	out, err := ploopOut("-v-1", "info", "-s", d.dd)
	if err == nil {
		i := reImageInfo.FindStringSubmatch(out)
		if len(i) != 4 {
			return info, &Err{c: -1,
				s: "Can't parse ploop info -s output:\n" + out}
		}
		info.Blocks, _ = strconv.ParseUint(i[1], 10, 64)
		bs, _ := strconv.ParseUint(i[2], 10, 32)
		info.BlockSize = uint32(bs)
		info.Version, _ = strconv.Atoi(i[3])
	}

	return info, err
}
<|endoftext|>


<commit_before>package transport

import (
	"bufio"
	"context"
	pb "github.com/Randomsock5/tcptunnel/proto"
	"log"
	"net"
	"time"

	"sync"
)

func handleErr(err error) {
	if err != nil {
		panic(err)
	}
}

func recoverHandle() {
	if rec := recover(); rec != nil {
		err := rec.(error)
		log.Fatalln(err)
	}
}

type proxyService struct {
	forward string
}

func (s *proxyService) Stream(stream pb.ProxyService_StreamServer) error {
	forwardConn, err := net.Dial("tcp", s.forward)
	if err != nil {
		log.Println(err)
		return err
	}
	defer forwardConn.Close()
	wg := &sync.WaitGroup{}

	go func() {
		defer wg.Done()
		defer recoverHandle()

		for {
			buf := make([]byte, 768)
			i, err := forwardConn.Read(buf)
			handleErr(err)

			var payload pb.Payload
			payload.Data = buf[:i]

			err = stream.Send(&payload)
			handleErr(err)
		}
	}()

	go func() {
		defer recoverHandle()
		defer wg.Done()
		writeBuff := bufio.NewWriter(forwardConn)

		for {
			payload, err := stream.Recv()
			handleErr(err)

			data := payload.GetData()
			_, err = writeBuff.Write(data)
			handleErr(err)

			err = writeBuff.Flush()
			handleErr(err)
		}
	}()

	wg.Add(2)
	wg.Wait()

	return nil
}

func NewServer(forward string) pb.ProxyServiceServer {
	s := &proxyService{forward: forward}
	return s
}

func ClientProxyService(conn net.Conn, client pb.ProxyServiceClient) {
	defer conn.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	stream, err := client.Stream(ctx)
	if err != nil {
		log.Println(err)
		return
	}
	defer stream.CloseSend()
	wg := &sync.WaitGroup{}

	go func() {
		defer recoverHandle()
		defer wg.Done()

		for {
			buf := make([]byte, 768)
			i, err := conn.Read(buf)
			handleErr(err)

			var payload pb.Payload
			payload.Data = buf[:i]

			err = stream.Send(&payload)
			handleErr(err)
		}
	}()

	go func() {
		defer recoverHandle()
		defer wg.Done()
		writeBuff := bufio.NewWriter(conn)

		for {
			payload, err := stream.Recv()
			handleErr(err)

			data := payload.GetData()
			_, err = writeBuff.Write(data)
			handleErr(err)

			err = writeBuff.Flush()
			handleErr(err)
		}
	}()

	wg.Add(2)
	wg.Wait()

	return
}
<commit_msg>fix client bug<commit_after>package transport

import (
	"bufio"
	"context"
	pb "github.com/Randomsock5/tcptunnel/proto"
	"log"
	"net"
	"time"

	"sync"
)

func handleErr(err error) {
	if err != nil {
		panic(err)
	}
}

func recoverHandle() {
	if rec := recover(); rec != nil {
		err := rec.(error)
		log.Println(err)
	}
}

type proxyService struct {
	forward string
}

func (s *proxyService) Stream(stream pb.ProxyService_StreamServer) error {
	forwardConn, err := net.Dial("tcp", s.forward)
	if err != nil {
		log.Println(err)
		return err
	}
	defer forwardConn.Close()
	wg := &sync.WaitGroup{}

	go func() {
		defer wg.Done()
		defer recoverHandle()

		for {
			buf := make([]byte, 768)
			i, err := forwardConn.Read(buf)
			handleErr(err)

			var payload pb.Payload
			payload.Data = buf[:i]

			err = stream.Send(&payload)
			handleErr(err)
		}
	}()

	go func() {
		defer recoverHandle()
		defer wg.Done()
		writeBuff := bufio.NewWriter(forwardConn)

		for {
			payload, err := stream.Recv()
			handleErr(err)

			data := payload.GetData()
			_, err = writeBuff.Write(data)
			handleErr(err)

			err = writeBuff.Flush()
			handleErr(err)
		}
	}()

	wg.Add(2)
	wg.Wait()

	return nil
}

func NewServer(forward string) pb.ProxyServiceServer {
	s := &proxyService{forward: forward}
	return s
}

func ClientProxyService(conn net.Conn, client pb.ProxyServiceClient) {
	defer conn.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	stream, err := client.Stream(ctx)
	if err != nil {
		log.Println(err)
		return
	}
	defer stream.CloseSend()
	wg := &sync.WaitGroup{}

	go func() {
		defer recoverHandle()
		defer wg.Done()

		for {
			buf := make([]byte, 768)
			i, err := conn.Read(buf)
			handleErr(err)

			var payload pb.Payload
			payload.Data = buf[:i]

			err = stream.Send(&payload)
			handleErr(err)
		}
	}()

	go func() {
		defer recoverHandle()
		defer wg.Done()
		writeBuff := bufio.NewWriter(conn)

		for {
			payload, err := stream.Recv()
			handleErr(err)

			data := payload.GetData()
			_, err = writeBuff.Write(data)
			handleErr(err)

			err = writeBuff.Flush()
			handleErr(err)
		}
	}()

	wg.Add(2)
	wg.Wait()

	return
}
<|endoftext|>


<commit_before>package htm

import (
	//"math"
	"bytes"
)

//Entries are positions of non-zero values
type SparseEntry struct {
	Row int
	Col int
}

//Sparse binary matrix stores indexes of non-zero entries in matrix
//to conserve space
type SparseBinaryMatrix struct {
	Width             int
	Height            int
	TotalNonZeroCount int
	Entries           []SparseEntry
}

//Create new sparse binary matrix of specified size
func NewSparseBinaryMatrix(height, width int) *SparseBinaryMatrix {
	m := &SparseBinaryMatrix{}
	m.Height = height
	m.Width = width
	//Intialize with 70% sparsity
	//m.Entries = make([]SparseEntry, int(math.Ceil(width*height*0.3)))
	return m
}

//Create sparse binary matrix from specified dense matrix
func NewSparseBinaryMatrixFromDense(values [][]bool) *SparseBinaryMatrix {
	if len(values) < 1 {
		panic("No values specified.")
	}
	m := &SparseBinaryMatrix{}
	m.Height = len(values)
	m.Width = len(values[0])

	for r := 0; r < m.Height; r++ {
		m.SetRowFromDense(r, values[r])
	}

	return m
}

// Creates a sparse binary matrix from specified integer array
// (any values greater than 0 are true)
func NewSparseBinaryMatrixFromInts(values [][]int) *SparseBinaryMatrix {
	if len(values) < 1 {
		panic("No values specified.")
	}

	m := &SparseBinaryMatrix{}
	m.Height = len(values)
	m.Width = len(values[0])

	for r := 0; r < m.Height; r++ {
		for c := 0; c < m.Width; c++ {
			if values[r][c] > 0 {
				m.Set(r, c, true)
			}
		}
	}

	return m
}

// func NewRandSparseBinaryMatrix() *SparseBinaryMatrix {
// }

// func (sm *SparseBinaryMatrix) Resize(width int, height int) {
// }

//Returns flattend dense represenation
func (sm *SparseBinaryMatrix) Flatten() []bool {
	result := make([]bool, sm.Height*sm.Width)
	for _, val := range sm.Entries {
		result[(val.Row*sm.Width)+val.Col] = true
	}
	return result
}

//Get value at col,row position
func (sm *SparseBinaryMatrix) Get(row int, col int) bool {
	for _, val := range sm.Entries {
		if val.Row == row && val.Col == col {
			return true
		}
	}
	return false
}

func (sm *SparseBinaryMatrix) delete(row int, col int) {
	for idx, val := range sm.Entries {
		if val.Row == row && val.Col == col {
			sm.Entries = append(sm.Entries[:idx], sm.Entries[idx+1:]...)
			break
		}
	}

}

//Set value at row,col position
func (sm *SparseBinaryMatrix) Set(row int, col int, value bool) {
	if !value {
		sm.delete(row, col)
		return
	}

	if sm.Get(row, col) {
		return
	}

	newEntry := SparseEntry{}
	newEntry.Col = col
	newEntry.Row = row
	sm.Entries = append(sm.Entries, newEntry)

}

//Replaces specified row with values, assumes values is ordered
//correctly
func (sm *SparseBinaryMatrix) ReplaceRow(row int, values []bool) {
	sm.validateRowCol(row, len(values))

	for i := 0; i < sm.Width; i++ {
		sm.Set(row, i, values[i])
	}
}

//Replaces row with true values at specified indices
func (sm *SparseBinaryMatrix) ReplaceRowByIndices(row int, indices []int) {
	sm.validateRow(row)

	for i := 0; i < sm.Width; i++ {
		val := false
		for x := 0; x < len(indices); x++ {
			if i == indices[x] {
				val = true
				break
			}
		}
		sm.Set(row, i, val)
	}
}

//Returns dense row
func (sm *SparseBinaryMatrix) GetDenseRow(row int) []bool {
	sm.validateRow(row)
	result := make([]bool, sm.Width)

	for i := 0; i < len(sm.Entries); i++ {
		if sm.Entries[i].Row == row {
			result[sm.Entries[i].Col] = true
		}
	}

	return result
}

//Returns a rows "on" indices
func (sm *SparseBinaryMatrix) GetRowIndices(row int) []int {
	result := []int{}
	for i := 0; i < len(sm.Entries); i++ {
		if sm.Entries[i].Row == row {
			result = append(result, sm.Entries[i].Col)
		}
	}
	return result
}

//Sets a sparse row from dense representation
func (sm *SparseBinaryMatrix) SetRowFromDense(row int, denseRow []bool) {
	sm.validateRowCol(row, len(denseRow))
	for i := 0; i < sm.Width; i++ {
		sm.Set(row, i, denseRow[i])
	}
}

//In a normal matrix this would be multiplication in binary terms
//we just and then sum the true entries
func (sm *SparseBinaryMatrix) RowAndSum(row []bool) []int {
	sm.validateCol(len(row))
	result := make([]int, sm.Height)

	for _, val := range sm.Entries {
		if row[val.Col] {
			result[val.Row]++
		}
	}

	return result
}

func (sm *SparseBinaryMatrix) ToString() string {
	var buffer bytes.Buffer

	for r := 0; r < sm.Height; r++ {
		for c := 0; c < sm.Width; c++ {
			if sm.Get(r, c) {
				buffer.WriteByte('1')
			} else {
				buffer.WriteByte('0')
			}
		}
		buffer.WriteByte('\n')
	}

	return buffer.String()
}

func (sm *SparseBinaryMatrix) validateCol(col int) {
	if col > sm.Width {
		panic("Specified row is wider than matrix.")
	}
}

func (sm *SparseBinaryMatrix) validateRow(row int) {
	if row > sm.Height {
		panic("Specified row is out of bounds.")
	}
}

func (sm *SparseBinaryMatrix) validateRowCol(row int, col int) {
	sm.validateCol(col)
	sm.validateRow(row)
}
<commit_msg>added clear method<commit_after>package htm

import (
	//"math"
	"bytes"
)

//Entries are positions of non-zero values
type SparseEntry struct {
	Row int
	Col int
}

//Sparse binary matrix stores indexes of non-zero entries in matrix
//to conserve space
type SparseBinaryMatrix struct {
	Width             int
	Height            int
	TotalNonZeroCount int
	Entries           []SparseEntry
}

//Create new sparse binary matrix of specified size
func NewSparseBinaryMatrix(height, width int) *SparseBinaryMatrix {
	m := &SparseBinaryMatrix{}
	m.Height = height
	m.Width = width
	//Intialize with 70% sparsity
	//m.Entries = make([]SparseEntry, int(math.Ceil(width*height*0.3)))
	return m
}

//Create sparse binary matrix from specified dense matrix
func NewSparseBinaryMatrixFromDense(values [][]bool) *SparseBinaryMatrix {
	if len(values) < 1 {
		panic("No values specified.")
	}
	m := &SparseBinaryMatrix{}
	m.Height = len(values)
	m.Width = len(values[0])

	for r := 0; r < m.Height; r++ {
		m.SetRowFromDense(r, values[r])
	}

	return m
}

// Creates a sparse binary matrix from specified integer array
// (any values greater than 0 are true)
func NewSparseBinaryMatrixFromInts(values [][]int) *SparseBinaryMatrix {
	if len(values) < 1 {
		panic("No values specified.")
	}

	m := &SparseBinaryMatrix{}
	m.Height = len(values)
	m.Width = len(values[0])

	for r := 0; r < m.Height; r++ {
		for c := 0; c < m.Width; c++ {
			if values[r][c] > 0 {
				m.Set(r, c, true)
			}
		}
	}

	return m
}

// func NewRandSparseBinaryMatrix() *SparseBinaryMatrix {
// }

// func (sm *SparseBinaryMatrix) Resize(width int, height int) {
// }

//Returns flattend dense represenation
func (sm *SparseBinaryMatrix) Flatten() []bool {
	result := make([]bool, sm.Height*sm.Width)
	for _, val := range sm.Entries {
		result[(val.Row*sm.Width)+val.Col] = true
	}
	return result
}

//Get value at col,row position
func (sm *SparseBinaryMatrix) Get(row int, col int) bool {
	for _, val := range sm.Entries {
		if val.Row == row && val.Col == col {
			return true
		}
	}
	return false
}

func (sm *SparseBinaryMatrix) delete(row int, col int) {
	for idx, val := range sm.Entries {
		if val.Row == row && val.Col == col {
			sm.Entries = append(sm.Entries[:idx], sm.Entries[idx+1:]...)
			break
		}
	}

}

//Set value at row,col position
func (sm *SparseBinaryMatrix) Set(row int, col int, value bool) {
	if !value {
		sm.delete(row, col)
		return
	}

	if sm.Get(row, col) {
		return
	}

	newEntry := SparseEntry{}
	newEntry.Col = col
	newEntry.Row = row
	sm.Entries = append(sm.Entries, newEntry)

}

//Replaces specified row with values, assumes values is ordered
//correctly
func (sm *SparseBinaryMatrix) ReplaceRow(row int, values []bool) {
	sm.validateRowCol(row, len(values))

	for i := 0; i < sm.Width; i++ {
		sm.Set(row, i, values[i])
	}
}

//Replaces row with true values at specified indices
func (sm *SparseBinaryMatrix) ReplaceRowByIndices(row int, indices []int) {
	sm.validateRow(row)

	for i := 0; i < sm.Width; i++ {
		val := false
		for x := 0; x < len(indices); x++ {
			if i == indices[x] {
				val = true
				break
			}
		}
		sm.Set(row, i, val)
	}
}

//Returns dense row
func (sm *SparseBinaryMatrix) GetDenseRow(row int) []bool {
	sm.validateRow(row)
	result := make([]bool, sm.Width)

	for i := 0; i < len(sm.Entries); i++ {
		if sm.Entries[i].Row == row {
			result[sm.Entries[i].Col] = true
		}
	}

	return result
}

//Returns a rows "on" indices
func (sm *SparseBinaryMatrix) GetRowIndices(row int) []int {
	result := []int{}
	for i := 0; i < len(sm.Entries); i++ {
		if sm.Entries[i].Row == row {
			result = append(result, sm.Entries[i].Col)
		}
	}
	return result
}

//Sets a sparse row from dense representation
func (sm *SparseBinaryMatrix) SetRowFromDense(row int, denseRow []bool) {
	sm.validateRowCol(row, len(denseRow))
	for i := 0; i < sm.Width; i++ {
		sm.Set(row, i, denseRow[i])
	}
}

//In a normal matrix this would be multiplication in binary terms
//we just and then sum the true entries
func (sm *SparseBinaryMatrix) RowAndSum(row []bool) []int {
	sm.validateCol(len(row))
	result := make([]int, sm.Height)

	for _, val := range sm.Entries {
		if row[val.Col] {
			result[val.Row]++
		}
	}

	return result
}

//Clears  all entries
func (sm *SparseBinaryMatrix) Clear() {
	sm.Entries = nil
}

func (sm *SparseBinaryMatrix) ToString() string {
	var buffer bytes.Buffer

	for r := 0; r < sm.Height; r++ {
		for c := 0; c < sm.Width; c++ {
			if sm.Get(r, c) {
				buffer.WriteByte('1')
			} else {
				buffer.WriteByte('0')
			}
		}
		buffer.WriteByte('\n')
	}

	return buffer.String()
}

func (sm *SparseBinaryMatrix) validateCol(col int) {
	if col > sm.Width {
		panic("Specified row is wider than matrix.")
	}
}

func (sm *SparseBinaryMatrix) validateRow(row int) {
	if row > sm.Height {
		panic("Specified row is out of bounds.")
	}
}

func (sm *SparseBinaryMatrix) validateRowCol(row int, col int) {
	sm.validateCol(col)
	sm.validateRow(row)
}
<|endoftext|>