<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	log "github.com/sirupsen/logrus"
)

func runVbox(args []string) {
	invoked := filepath.Base(os.Args[0])
	flags := flag.NewFlagSet("vbox", flag.ExitOnError)
	flags.Usage = func() {
		fmt.Printf("USAGE: %s run vbox [options] path\n\n", invoked)
		fmt.Printf("'path' specifies the path to the VM image.\n")
		fmt.Printf("\n")
		fmt.Printf("Options:\n")
		flags.PrintDefaults()
		fmt.Printf("\n")
	}

	// Display flags
	enableGUI := flags.Bool("gui", false, "Show the VM GUI")

	// vbox options
	vboxmanageFlag := flags.String("vboxmanage", "VBoxManage", "VBoxManage binary to use")
	keep := flags.Bool("keep", false, "Keep the VM after finishing")
	vmName := flags.String("name", "", "Name of the Virtualbox VM")
	state := flags.String("state", "", "Path to directory to keep VM state in")

	// Paths and settings for disks
	var disks Disks
	flags.Var(&disks, "disk", "Disk config, may be repeated. [file=]path[,size=1G][,format=raw]")

	// VM configuration
	cpus := flags.String("cpus", "1", "Number of CPUs")
	mem := flags.String("mem", "1024", "Amount of memory in MB")

	// booting config
	isoBoot := flags.Bool("iso", false, "Boot image is an ISO")
	uefiBoot := flags.Bool("uefi", false, "Use UEFI boot")

	// networking
	networking := flags.String("networking", "nat", "Networking mode. null|nat|bridged|intnet|hostonly|generic|natnetwork[<devicename>]")
	bridgeadapter := flags.String("bridgeadapter", "", "Bridge adapter interface to use if networking mode is bridged")

	if err := flags.Parse(args); err != nil {
		log.Fatal("Unable to parse args")
	}
	remArgs := flags.Args()

	if runtime.GOOS == "windows" {
		log.Fatalf("TODO: Windows is not yet supported")
	}

	if len(remArgs) == 0 {
		fmt.Println("Please specify the path to the image to boot")
		flags.Usage()
		os.Exit(1)
	}
	path := remArgs[0]

	if strings.HasSuffix(path, ".iso") {
		*isoBoot = true
	}

	vboxmanage, err := exec.LookPath(*vboxmanageFlag)
	if err != nil {
		log.Fatalf("Cannot find management binary %s: %v", *vboxmanageFlag, err)
	}

	name := *vmName
	if name == "" {
		name = strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
	}

	if *state == "" {
		prefix := strings.TrimSuffix(path, filepath.Ext(path))
		*state = prefix + "-state"
	}
	if err := os.MkdirAll(*state, 0755); err != nil {
		log.Fatalf("Could not create state directory: %v", err)
	}

	// remove machine in case it already exists
	cleanup(vboxmanage, name, false)

	_, out, err := manage(vboxmanage, "createvm", "--name", name, "--register")
	if err != nil {
		log.Fatalf("createvm error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--acpi", "on")
	if err != nil {
		log.Fatalf("modifyvm --acpi error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--memory", *mem)
	if err != nil {
		log.Fatalf("modifyvm --memory error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--cpus", *cpus)
	if err != nil {
		log.Fatalf("modifyvm --cpus error: %v\n%s", err, out)
	}

	firmware := "bios"
	if *uefiBoot {
		firmware = "efi"
	}
	_, out, err = manage(vboxmanage, "modifyvm", name, "--firmware", firmware)
	if err != nil {
		log.Fatalf("modifyvm --firmware error: %v\n%s", err, out)
	}

	// set up serial console
	_, out, err = manage(vboxmanage, "modifyvm", name, "--uart1", "0x3F8", "4")
	if err != nil {
		log.Fatalf("modifyvm --uart error: %v\n%s", err, out)
	}

	var consolePath string
	if runtime.GOOS == "windows" {
		// TODO use a named pipe on Windows
	} else {
		consolePath = filepath.Join(*state, "console")
		consolePath, err = filepath.Abs(consolePath)
		if err != nil {
			log.Fatalf("Bad path: %v", err)
		}
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--uartmode1", "client", consolePath)
	if err != nil {
		log.Fatalf("modifyvm --uartmode error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "storagectl", name, "--name", "IDE Controller", "--add", "ide")
	if err != nil {
		log.Fatalf("storagectl error: %v\n%s", err, out)
	}

	if *isoBoot {
		_, out, err = manage(vboxmanage, "storageattach", name, "--storagectl", "IDE Controller", "--port", "1", "--device", "0", "--type", "dvddrive", "--medium", path)
		if err != nil {
			log.Fatalf("storageattach error: %v\n%s", err, out)
		}
		_, out, err = manage(vboxmanage, "modifyvm", name, "--boot1", "dvd")
		if err != nil {
			log.Fatalf("modifyvm --boot error: %v\n%s", err, out)
		}
	} else {
		_, out, err = manage(vboxmanage, "storageattach", name, "--storagectl", "IDE Controller", "--port", "1", "--device", "0", "--type", "hdd", "--medium", path)
		if err != nil {
			log.Fatalf("storageattach error: %v\n%s", err, out)
		}
		_, out, err = manage(vboxmanage, "modifyvm", name, "--boot1", "disk")
		if err != nil {
			log.Fatalf("modifyvm --boot error: %v\n%s", err, out)
		}
	}

	for i, d := range disks {
		id := strconv.Itoa(i)
		if d.Size != 0 && d.Format == "" {
			d.Format = "raw"
		}
		if d.Format != "raw" && d.Path == "" {
			log.Fatal("vbox currently can only create raw disks")
		}
		if d.Path == "" && d.Size == 0 {
			log.Fatal("please specify an existing disk file or a size")
		}
		if d.Path == "" {
			d.Path = filepath.Join(*state, "disk"+id+".img")
			if err := os.Truncate(d.Path, int64(d.Size)*int64(1048576)); err != nil {
				log.Fatalf("Cannot create disk: %v", err)
			}
		}
		_, out, err = manage(vboxmanage, "storageattach", name, "--storagectl", "IDE Controller", "--port", "2", "--device", id, "--type", "hdd", "--medium", d.Path)
		if err != nil {
			log.Fatalf("storageattach error: %v\n%s", err, out)
		}
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--nictype1", "virtio")
	if err != nil {
		log.Fatalf("modifyvm --nictype error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--nic1", *networking)
	if err != nil {
		log.Fatalf("modifyvm --nic error: %v\n%s", err, out)
	}
	if *networking == "bridged" {
		_, out, err = manage(vboxmanage, "modifyvm", name, "--bridgeadapter1", *bridgeadapter)
		if err != nil {
			log.Fatalf("modifyvm --bridgeadapter error: %v\n%s", err, out)
		}
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--cableconnected1", "on")
	if err != nil {
		log.Fatalf("modifyvm --cableconnected error: %v\n%s", err, out)
	}

	// create socket
	_ = os.Remove(consolePath)
	ln, err := net.Listen("unix", consolePath)
	if err != nil {
		log.Fatalf("Cannot listen on console socket %s: %v", consolePath, err)
	}

	var vmType string
	if *enableGUI {
		vmType = "gui"
	} else {
		vmType = "headless"
	}

	_, out, err = manage(vboxmanage, "startvm", name, "--type", vmType)
	if err != nil {
		log.Fatalf("startvm error: %v\n%s", err, out)
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c
		cleanup(vboxmanage, name, *keep)
		os.Exit(1)
	}()

	socket, err := ln.Accept()
	if err != nil {
		log.Fatalf("Accept error: %v", err)
	}

	go func() {
		if _, err := io.Copy(socket, os.Stdin); err != nil {
			cleanup(vboxmanage, name, *keep)
			log.Fatalf("Copy error: %v", err)
		}
		cleanup(vboxmanage, name, *keep)
		os.Exit(0)
	}()
	go func() {
		if _, err := io.Copy(os.Stdout, socket); err != nil {
			cleanup(vboxmanage, name, *keep)
			log.Fatalf("Copy error: %v", err)
		}
		cleanup(vboxmanage, name, *keep)
		os.Exit(0)
	}()
	// wait forever
	select {}
}

func cleanup(vboxmanage string, name string, keep bool) {
	_, _, _ = manage(vboxmanage, "controlvm", name, "poweroff")

	if keep {
		return
	}

	// delete VM
	_, _, _ = manage(vboxmanage, "unregistervm", name, "--delete")
}

func manage(vboxmanage string, args ...string) (string, string, error) {
	cmd := exec.Command(vboxmanage, args...)
	log.Debugf("[VBOX]: %s %s", vboxmanage, strings.Join(args, " "))

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	return stdout.String(), stderr.String(), err
}
<commit_msg>Allow `linuxkit run vbox` to use multiple drives<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"

	log "github.com/sirupsen/logrus"
)

func runVbox(args []string) {
	invoked := filepath.Base(os.Args[0])
	flags := flag.NewFlagSet("vbox", flag.ExitOnError)
	flags.Usage = func() {
		fmt.Printf("USAGE: %s run vbox [options] path\n\n", invoked)
		fmt.Printf("'path' specifies the path to the VM image.\n")
		fmt.Printf("\n")
		fmt.Printf("Options:\n")
		flags.PrintDefaults()
		fmt.Printf("\n")
	}

	// Display flags
	enableGUI := flags.Bool("gui", false, "Show the VM GUI")

	// vbox options
	vboxmanageFlag := flags.String("vboxmanage", "VBoxManage", "VBoxManage binary to use")
	keep := flags.Bool("keep", false, "Keep the VM after finishing")
	vmName := flags.String("name", "", "Name of the Virtualbox VM")
	state := flags.String("state", "", "Path to directory to keep VM state in")

	// Paths and settings for disks
	var disks Disks
	flags.Var(&disks, "disk", "Disk config, may be repeated. [file=]path[,size=1G][,format=raw]")

	// VM configuration
	cpus := flags.String("cpus", "1", "Number of CPUs")
	mem := flags.String("mem", "1024", "Amount of memory in MB")

	// booting config
	isoBoot := flags.Bool("iso", false, "Boot image is an ISO")
	uefiBoot := flags.Bool("uefi", false, "Use UEFI boot")

	// networking
	networking := flags.String("networking", "nat", "Networking mode. null|nat|bridged|intnet|hostonly|generic|natnetwork[<devicename>]")
	bridgeadapter := flags.String("bridgeadapter", "", "Bridge adapter interface to use if networking mode is bridged")

	if err := flags.Parse(args); err != nil {
		log.Fatal("Unable to parse args")
	}
	remArgs := flags.Args()

	if runtime.GOOS == "windows" {
		log.Fatalf("TODO: Windows is not yet supported")
	}

	if len(remArgs) == 0 {
		fmt.Println("Please specify the path to the image to boot")
		flags.Usage()
		os.Exit(1)
	}
	path := remArgs[0]

	if strings.HasSuffix(path, ".iso") {
		*isoBoot = true
	}

	vboxmanage, err := exec.LookPath(*vboxmanageFlag)
	if err != nil {
		log.Fatalf("Cannot find management binary %s: %v", *vboxmanageFlag, err)
	}

	name := *vmName
	if name == "" {
		name = strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
	}

	if *state == "" {
		prefix := strings.TrimSuffix(path, filepath.Ext(path))
		*state = prefix + "-state"
	}
	if err := os.MkdirAll(*state, 0755); err != nil {
		log.Fatalf("Could not create state directory: %v", err)
	}

	// remove machine in case it already exists
	cleanup(vboxmanage, name, false)

	_, out, err := manage(vboxmanage, "createvm", "--name", name, "--register")
	if err != nil {
		log.Fatalf("createvm error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--acpi", "on")
	if err != nil {
		log.Fatalf("modifyvm --acpi error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--memory", *mem)
	if err != nil {
		log.Fatalf("modifyvm --memory error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--cpus", *cpus)
	if err != nil {
		log.Fatalf("modifyvm --cpus error: %v\n%s", err, out)
	}

	firmware := "bios"
	if *uefiBoot {
		firmware = "efi"
	}
	_, out, err = manage(vboxmanage, "modifyvm", name, "--firmware", firmware)
	if err != nil {
		log.Fatalf("modifyvm --firmware error: %v\n%s", err, out)
	}

	// set up serial console
	_, out, err = manage(vboxmanage, "modifyvm", name, "--uart1", "0x3F8", "4")
	if err != nil {
		log.Fatalf("modifyvm --uart error: %v\n%s", err, out)
	}

	var consolePath string
	if runtime.GOOS == "windows" {
		// TODO use a named pipe on Windows
	} else {
		consolePath = filepath.Join(*state, "console")
		consolePath, err = filepath.Abs(consolePath)
		if err != nil {
			log.Fatalf("Bad path: %v", err)
		}
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--uartmode1", "client", consolePath)
	if err != nil {
		log.Fatalf("modifyvm --uartmode error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "storagectl", name, "--name", "IDE Controller", "--add", "ide")
	if err != nil {
		log.Fatalf("storagectl error: %v\n%s", err, out)
	}

	if *isoBoot {
		_, out, err = manage(vboxmanage, "storageattach", name, "--storagectl", "IDE Controller", "--port", "1", "--device", "0", "--type", "dvddrive", "--medium", path)
		if err != nil {
			log.Fatalf("storageattach error: %v\n%s", err, out)
		}
		_, out, err = manage(vboxmanage, "modifyvm", name, "--boot1", "dvd")
		if err != nil {
			log.Fatalf("modifyvm --boot error: %v\n%s", err, out)
		}
	} else {
		_, out, err = manage(vboxmanage, "storageattach", name, "--storagectl", "IDE Controller", "--port", "1", "--device", "0", "--type", "hdd", "--medium", path)
		if err != nil {
			log.Fatalf("storageattach error: %v\n%s", err, out)
		}
		_, out, err = manage(vboxmanage, "modifyvm", name, "--boot1", "disk")
		if err != nil {
			log.Fatalf("modifyvm --boot error: %v\n%s", err, out)
		}
	}

	if len(disks) > 0 {
		_, out, err = manage(vboxmanage, "storagectl", name, "--name", "SATA", "--add", "sata")
		if err != nil {
			log.Fatalf("storagectl error: %v\n%s", err, out)
		}
	}

	for i, d := range disks {
		id := strconv.Itoa(i)
		if d.Size != 0 && d.Format == "" {
			d.Format = "raw"
		}
		if d.Format != "raw" && d.Path == "" {
			log.Fatal("vbox currently can only create raw disks")
		}
		if d.Path == "" && d.Size == 0 {
			log.Fatal("please specify an existing disk file or a size")
		}
		if d.Path == "" {
			d.Path = filepath.Join(*state, "disk"+id+".img")
			if err := os.Truncate(d.Path, int64(d.Size)*int64(1048576)); err != nil {
				log.Fatalf("Cannot create disk: %v", err)
			}
		}
		_, out, err = manage(vboxmanage, "storageattach", name, "--storagectl", "SATA", "--port", "0", "--device", id, "--type", "hdd", "--medium", d.Path)
		if err != nil {
			log.Fatalf("storageattach error: %v\n%s", err, out)
		}
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--nictype1", "virtio")
	if err != nil {
		log.Fatalf("modifyvm --nictype error: %v\n%s", err, out)
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--nic1", *networking)
	if err != nil {
		log.Fatalf("modifyvm --nic error: %v\n%s", err, out)
	}
	if *networking == "bridged" {
		_, out, err = manage(vboxmanage, "modifyvm", name, "--bridgeadapter1", *bridgeadapter)
		if err != nil {
			log.Fatalf("modifyvm --bridgeadapter error: %v\n%s", err, out)
		}
	}

	_, out, err = manage(vboxmanage, "modifyvm", name, "--cableconnected1", "on")
	if err != nil {
		log.Fatalf("modifyvm --cableconnected error: %v\n%s", err, out)
	}

	// create socket
	_ = os.Remove(consolePath)
	ln, err := net.Listen("unix", consolePath)
	if err != nil {
		log.Fatalf("Cannot listen on console socket %s: %v", consolePath, err)
	}

	var vmType string
	if *enableGUI {
		vmType = "gui"
	} else {
		vmType = "headless"
	}

	_, out, err = manage(vboxmanage, "startvm", name, "--type", vmType)
	if err != nil {
		log.Fatalf("startvm error: %v\n%s", err, out)
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		<-c
		cleanup(vboxmanage, name, *keep)
		os.Exit(1)
	}()

	socket, err := ln.Accept()
	if err != nil {
		log.Fatalf("Accept error: %v", err)
	}

	go func() {
		if _, err := io.Copy(socket, os.Stdin); err != nil {
			cleanup(vboxmanage, name, *keep)
			log.Fatalf("Copy error: %v", err)
		}
		cleanup(vboxmanage, name, *keep)
		os.Exit(0)
	}()
	go func() {
		if _, err := io.Copy(os.Stdout, socket); err != nil {
			cleanup(vboxmanage, name, *keep)
			log.Fatalf("Copy error: %v", err)
		}
		cleanup(vboxmanage, name, *keep)
		os.Exit(0)
	}()
	// wait forever
	select {}
}

func cleanup(vboxmanage string, name string, keep bool) {
	_, _, _ = manage(vboxmanage, "controlvm", name, "poweroff")

	if keep {
		return
	}

	// delete VM
	_, _, _ = manage(vboxmanage, "unregistervm", name, "--delete")
}

func manage(vboxmanage string, args ...string) (string, string, error) {
	cmd := exec.Command(vboxmanage, args...)
	log.Debugf("[VBOX]: %s %s", vboxmanage, strings.Join(args, " "))

	var stdout bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()
	return stdout.String(), stderr.String(), err
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/hyperkit/go"
)

// Process the run arguments and execute run
func runHyperKit(args []string) {
	hyperkitCmd := flag.NewFlagSet("hyperkit", flag.ExitOnError)
	hyperkitCmd.Usage = func() {
		fmt.Printf("USAGE: %s run hyperkit [options] [prefix]\n\n", os.Args[0])
		fmt.Printf("'prefix' specifies the path to the VM image.\n")
		fmt.Printf("It defaults to './moby'.\n")
		fmt.Printf("\n")
		fmt.Printf("Options:\n")
		hyperkitCmd.PrintDefaults()
	}
	runCPUs := hyperkitCmd.Int("cpus", 1, "Number of CPUs")
	runMem := hyperkitCmd.Int("mem", 1024, "Amount of memory in MB")
	runDiskSz := hyperkitCmd.Int("disk-size", 0, "Size of Disk in MB")
	runDisk := hyperkitCmd.String("disk", "", "Path to disk image to used")

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

	prefix := "moby"
	if len(remArgs) > 0 {
		prefix = remArgs[0]
	}

	runHyperKitInternal(*runCPUs, *runMem, *runDiskSz, *runDisk, prefix)
}

func runHyperKitInternal(cpus, mem, diskSz int, disk, prefix string) {
	cmdline, err := ioutil.ReadFile(prefix + "-cmdline")
	if err != nil {
		log.Fatalf("Cannot open cmdline file: %v", err)
	}

	if diskSz != 0 && disk == "" {
		disk = prefix + "-disk.img"
	}

	h, err := hyperkit.New("", "", "auto", disk)
	if err != nil {
		log.Fatalln("Error creating hyperkit: ", err)
	}

	h.Kernel = prefix + "-bzImage"
	h.Initrd = prefix + "-initrd.img"
	h.CPUs = cpus
	h.Memory = mem
	h.DiskSize = diskSz

	err = h.Run(string(cmdline))
	if err != nil {
		log.Fatalf("Cannot run hyperkit: %v", err)
	}
}
<commit_msg>cli: Add option to specify hyperkit to use<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/hyperkit/go"
)

// Process the run arguments and execute run
func runHyperKit(args []string) {
	hyperkitCmd := flag.NewFlagSet("hyperkit", flag.ExitOnError)
	hyperkitCmd.Usage = func() {
		fmt.Printf("USAGE: %s run hyperkit [options] [prefix]\n\n", os.Args[0])
		fmt.Printf("'prefix' specifies the path to the VM image.\n")
		fmt.Printf("It defaults to './moby'.\n")
		fmt.Printf("\n")
		fmt.Printf("Options:\n")
		hyperkitCmd.PrintDefaults()
	}
	runHyperKit := hyperkitCmd.String("hyperkit", "", "Path to hyperkit binary (if not in default location)")
	runCPUs := hyperkitCmd.Int("cpus", 1, "Number of CPUs")
	runMem := hyperkitCmd.Int("mem", 1024, "Amount of memory in MB")
	runDiskSz := hyperkitCmd.Int("disk-size", 0, "Size of Disk in MB")
	runDisk := hyperkitCmd.String("disk", "", "Path to disk image to used")

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

	prefix := "moby"
	if len(remArgs) > 0 {
		prefix = remArgs[0]
	}

	runHyperKitInternal(*runHyperKit, *runCPUs, *runMem, *runDiskSz, *runDisk, prefix)
}

func runHyperKitInternal(hyperkitPath string, cpus, mem, diskSz int, disk, prefix string) {
	cmdline, err := ioutil.ReadFile(prefix + "-cmdline")
	if err != nil {
		log.Fatalf("Cannot open cmdline file: %v", err)
	}

	if diskSz != 0 && disk == "" {
		disk = prefix + "-disk.img"
	}

	h, err := hyperkit.New(hyperkitPath, "", "auto", disk)
	if err != nil {
		log.Fatalln("Error creating hyperkit: ", err)
	}

	h.Kernel = prefix + "-bzImage"
	h.Initrd = prefix + "-initrd.img"
	h.CPUs = cpus
	h.Memory = mem
	h.DiskSize = diskSz

	err = h.Run(string(cmdline))
	if err != nil {
		log.Fatalf("Cannot run hyperkit: %v", err)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2013 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
//-----------------------------------------------------------------------------

package mop

import (
	`sort`
	`strings`
	`strconv`
)

type sortable []Stock
func (list sortable) Len() int { return len(list) }
func (list sortable) Swap(i, j int) { list[i], list[j] = list[j], list[i] }

type byTickerAsc       struct { sortable }
type byLastTradeAsc    struct { sortable }
type byChangeAsc       struct { sortable }
type byChangePctAsc    struct { sortable }
type byOpenAsc         struct { sortable }
type byLowAsc          struct { sortable }
type byHighAsc         struct { sortable }
type byLow52Asc        struct { sortable }
type byHigh52Asc       struct { sortable }
type byVolumeAsc       struct { sortable }
type byAvgVolumeAsc    struct { sortable }
type byPeRatioAsc      struct { sortable }
type byDividendAsc     struct { sortable }
type byYieldAsc        struct { sortable }
type byMarketCapAsc    struct { sortable }

type byTickerDesc      struct { sortable }
type byLastTradeDesc   struct { sortable }
type byChangeDesc      struct { sortable }
type byChangePctDesc   struct { sortable }
type byOpenDesc        struct { sortable }
type byLowDesc         struct { sortable }
type byHighDesc        struct { sortable }
type byLow52Desc       struct { sortable }
type byHigh52Desc      struct { sortable }
type byVolumeDesc      struct { sortable }
type byAvgVolumeDesc   struct { sortable }
type byPeRatioDesc     struct { sortable }
type byDividendDesc    struct { sortable }
type byYieldDesc       struct { sortable }
type byMarketCapDesc   struct { sortable }

func (list byTickerAsc)       Less(i, j int) bool { return list.sortable[i].Ticker        < list.sortable[j].Ticker }
func (list byLastTradeAsc)    Less(i, j int) bool { return list.sortable[i].LastTrade     < list.sortable[j].LastTrade }
func (list byChangeAsc)       Less(i, j int) bool { return c(list.sortable[i].Change)     < c(list.sortable[j].Change) }
func (list byChangePctAsc)    Less(i, j int) bool { return c(list.sortable[i].ChangePct)  < c(list.sortable[j].ChangePct) }
func (list byOpenAsc)         Less(i, j int) bool { return list.sortable[i].Open          < list.sortable[j].Open }
func (list byLowAsc)          Less(i, j int) bool { return list.sortable[i].Low           < list.sortable[j].Low }
func (list byHighAsc)         Less(i, j int) bool { return list.sortable[i].High          < list.sortable[j].High }
func (list byLow52Asc)        Less(i, j int) bool { return list.sortable[i].Low52         < list.sortable[j].Low52 }
func (list byHigh52Asc)       Less(i, j int) bool { return list.sortable[i].High52        < list.sortable[j].High52 }
func (list byVolumeAsc)       Less(i, j int) bool { return list.sortable[i].Volume        < list.sortable[j].Volume }
func (list byAvgVolumeAsc)    Less(i, j int) bool { return list.sortable[i].AvgVolume     < list.sortable[j].AvgVolume }
func (list byPeRatioAsc)      Less(i, j int) bool { return list.sortable[i].PeRatio       < list.sortable[j].PeRatio }
func (list byDividendAsc)     Less(i, j int) bool { return list.sortable[i].Dividend      < list.sortable[j].Dividend }
func (list byYieldAsc)        Less(i, j int) bool { return list.sortable[i].Yield         < list.sortable[j].Yield }
func (list byMarketCapAsc)    Less(i, j int) bool { return m(list.sortable[i].MarketCap)  < m(list.sortable[j].MarketCap) }

func (list byTickerDesc)      Less(i, j int) bool { return list.sortable[j].Ticker        < list.sortable[i].Ticker }
func (list byLastTradeDesc)   Less(i, j int) bool { return list.sortable[j].LastTrade     < list.sortable[i].LastTrade }
func (list byChangeDesc)      Less(i, j int) bool { return c(list.sortable[j].ChangePct)  < c(list.sortable[i].ChangePct) }
func (list byChangePctDesc)   Less(i, j int) bool { return c(list.sortable[j].ChangePct)  < c(list.sortable[i].ChangePct) }
func (list byOpenDesc)        Less(i, j int) bool { return list.sortable[j].Open          < list.sortable[i].Open }
func (list byLowDesc)         Less(i, j int) bool { return list.sortable[j].Low           < list.sortable[i].Low }
func (list byHighDesc)        Less(i, j int) bool { return list.sortable[j].High          < list.sortable[i].High }
func (list byLow52Desc)       Less(i, j int) bool { return list.sortable[j].Low52         < list.sortable[i].Low52 }
func (list byHigh52Desc)      Less(i, j int) bool { return list.sortable[j].High52        < list.sortable[i].High52 }
func (list byVolumeDesc)      Less(i, j int) bool { return list.sortable[j].Volume        < list.sortable[i].Volume }
func (list byAvgVolumeDesc)   Less(i, j int) bool { return list.sortable[j].AvgVolume     < list.sortable[i].AvgVolume }
func (list byPeRatioDesc)     Less(i, j int) bool { return list.sortable[j].PeRatio       < list.sortable[i].PeRatio }
func (list byDividendDesc)    Less(i, j int) bool { return list.sortable[j].Dividend      < list.sortable[i].Dividend }
func (list byYieldDesc)       Less(i, j int) bool { return list.sortable[j].Yield         < list.sortable[i].Yield }
func (list byMarketCapDesc)   Less(i, j int) bool { return m(list.sortable[j].MarketCap)  < m(list.sortable[i].MarketCap) }

type Sorter struct {
	profile  *Profile
}

func (self *Sorter) Initialize(profile *Profile) *Sorter {
	self.profile = profile

	return self
}

func (self *Sorter) SortByCurrentColumn(stocks []Stock) *Sorter {
	var interfaces []sort.Interface

	if self.profile.Ascending {
		interfaces = []sort.Interface{
			byTickerAsc       { stocks },
			byLastTradeAsc    { stocks },
			byChangeAsc       { stocks },
			byChangePctAsc    { stocks },
			byOpenAsc         { stocks },
			byLowAsc          { stocks },
			byHighAsc         { stocks },
			byLow52Asc        { stocks },
			byHigh52Asc       { stocks },
			byVolumeAsc       { stocks },
			byAvgVolumeAsc    { stocks },
			byPeRatioAsc      { stocks },
			byDividendAsc     { stocks },
			byYieldAsc        { stocks },
			byMarketCapAsc    { stocks },
		}
	} else {
		interfaces = []sort.Interface{
			byTickerDesc      { stocks },
			byLastTradeDesc   { stocks },
			byChangeDesc      { stocks },
			byChangePctDesc   { stocks },
			byOpenDesc        { stocks },
			byLowDesc         { stocks },
			byHighDesc        { stocks },
			byLow52Desc       { stocks },
			byHigh52Desc      { stocks },
			byVolumeDesc      { stocks },
			byAvgVolumeDesc   { stocks },
			byPeRatioDesc     { stocks },
			byDividendDesc    { stocks },
			byYieldDesc       { stocks },
			byMarketCapDesc   { stocks },
		}
	}

	sort.Sort(interfaces[self.profile.SortColumn])

	return self
}

// The same exact method is used to sort by Change and Change%. In both cases
// we sort by the value of Change% so that $0.00 change gets sorted proferly.
func c(str string) float32 {
	trimmed := strings.Replace(strings.Trim(str, ` %`), `$`, ``, 1)
	value, _ := strconv.ParseFloat(trimmed, 32)
	return float32(value)
}

func m(str string) float32 {
	multiplier := 1.0
	switch str[len(str)-1:len(str)] {
	case `B`:
		multiplier = 1000000000.0
	case `M`:
		multiplier = 1000000.0
	case `K`:
		multiplier = 1000.0
	}
	trimmed := strings.Trim(str, ` $BMK`)
	value, _ := strconv.ParseFloat(trimmed, 32)
	return float32(value * multiplier)
}
<commit_msg>Documented the sorter<commit_after>// Copyright (c) 2013 by Michael Dvorkin. All Rights Reserved.
// Use of this source code is governed by a MIT-style license that can
// be found in the LICENSE file.

package mop

import (
	`sort`
	`strconv`
	`strings`
)

// Sorter gets called to sort stock quotes by one of the columns. The
// setup is rather lengthy; there should probably be more concise way
// that uses reflection and avoids hardcoding the column names.
type Sorter struct {
	profile  *Profile  // Pointer to where we store sort column and order.
}

type sortable []Stock
func (list sortable) Len() int { return len(list) }
func (list sortable) Swap(i, j int) { list[i], list[j] = list[j], list[i] }

type byTickerAsc       struct { sortable }
type byLastTradeAsc    struct { sortable }
type byChangeAsc       struct { sortable }
type byChangePctAsc    struct { sortable }
type byOpenAsc         struct { sortable }
type byLowAsc          struct { sortable }
type byHighAsc         struct { sortable }
type byLow52Asc        struct { sortable }
type byHigh52Asc       struct { sortable }
type byVolumeAsc       struct { sortable }
type byAvgVolumeAsc    struct { sortable }
type byPeRatioAsc      struct { sortable }
type byDividendAsc     struct { sortable }
type byYieldAsc        struct { sortable }
type byMarketCapAsc    struct { sortable }

type byTickerDesc      struct { sortable }
type byLastTradeDesc   struct { sortable }
type byChangeDesc      struct { sortable }
type byChangePctDesc   struct { sortable }
type byOpenDesc        struct { sortable }
type byLowDesc         struct { sortable }
type byHighDesc        struct { sortable }
type byLow52Desc       struct { sortable }
type byHigh52Desc      struct { sortable }
type byVolumeDesc      struct { sortable }
type byAvgVolumeDesc   struct { sortable }
type byPeRatioDesc     struct { sortable }
type byDividendDesc    struct { sortable }
type byYieldDesc       struct { sortable }
type byMarketCapDesc   struct { sortable }

func (list byTickerAsc)       Less(i, j int) bool { return list.sortable[i].Ticker        < list.sortable[j].Ticker }
func (list byLastTradeAsc)    Less(i, j int) bool { return list.sortable[i].LastTrade     < list.sortable[j].LastTrade }
func (list byChangeAsc)       Less(i, j int) bool { return c(list.sortable[i].Change)     < c(list.sortable[j].Change) }
func (list byChangePctAsc)    Less(i, j int) bool { return c(list.sortable[i].ChangePct)  < c(list.sortable[j].ChangePct) }
func (list byOpenAsc)         Less(i, j int) bool { return list.sortable[i].Open          < list.sortable[j].Open }
func (list byLowAsc)          Less(i, j int) bool { return list.sortable[i].Low           < list.sortable[j].Low }
func (list byHighAsc)         Less(i, j int) bool { return list.sortable[i].High          < list.sortable[j].High }
func (list byLow52Asc)        Less(i, j int) bool { return list.sortable[i].Low52         < list.sortable[j].Low52 }
func (list byHigh52Asc)       Less(i, j int) bool { return list.sortable[i].High52        < list.sortable[j].High52 }
func (list byVolumeAsc)       Less(i, j int) bool { return list.sortable[i].Volume        < list.sortable[j].Volume }
func (list byAvgVolumeAsc)    Less(i, j int) bool { return list.sortable[i].AvgVolume     < list.sortable[j].AvgVolume }
func (list byPeRatioAsc)      Less(i, j int) bool { return list.sortable[i].PeRatio       < list.sortable[j].PeRatio }
func (list byDividendAsc)     Less(i, j int) bool { return list.sortable[i].Dividend      < list.sortable[j].Dividend }
func (list byYieldAsc)        Less(i, j int) bool { return list.sortable[i].Yield         < list.sortable[j].Yield }
func (list byMarketCapAsc)    Less(i, j int) bool { return m(list.sortable[i].MarketCap)  < m(list.sortable[j].MarketCap) }

func (list byTickerDesc)      Less(i, j int) bool { return list.sortable[j].Ticker        < list.sortable[i].Ticker }
func (list byLastTradeDesc)   Less(i, j int) bool { return list.sortable[j].LastTrade     < list.sortable[i].LastTrade }
func (list byChangeDesc)      Less(i, j int) bool { return c(list.sortable[j].ChangePct)  < c(list.sortable[i].ChangePct) }
func (list byChangePctDesc)   Less(i, j int) bool { return c(list.sortable[j].ChangePct)  < c(list.sortable[i].ChangePct) }
func (list byOpenDesc)        Less(i, j int) bool { return list.sortable[j].Open          < list.sortable[i].Open }
func (list byLowDesc)         Less(i, j int) bool { return list.sortable[j].Low           < list.sortable[i].Low }
func (list byHighDesc)        Less(i, j int) bool { return list.sortable[j].High          < list.sortable[i].High }
func (list byLow52Desc)       Less(i, j int) bool { return list.sortable[j].Low52         < list.sortable[i].Low52 }
func (list byHigh52Desc)      Less(i, j int) bool { return list.sortable[j].High52        < list.sortable[i].High52 }
func (list byVolumeDesc)      Less(i, j int) bool { return list.sortable[j].Volume        < list.sortable[i].Volume }
func (list byAvgVolumeDesc)   Less(i, j int) bool { return list.sortable[j].AvgVolume     < list.sortable[i].AvgVolume }
func (list byPeRatioDesc)     Less(i, j int) bool { return list.sortable[j].PeRatio       < list.sortable[i].PeRatio }
func (list byDividendDesc)    Less(i, j int) bool { return list.sortable[j].Dividend      < list.sortable[i].Dividend }
func (list byYieldDesc)       Less(i, j int) bool { return list.sortable[j].Yield         < list.sortable[i].Yield }
func (list byMarketCapDesc)   Less(i, j int) bool { return m(list.sortable[j].MarketCap)  < m(list.sortable[i].MarketCap) }

// Initialize simply saves the pointer to Profile for later use.
func (sorter *Sorter) Initialize(profile *Profile) *Sorter {
	sorter.profile = profile

	return sorter
}

// SortByCurrentColumn builds a list of sort interface based on current sort
// order, then calls sort.Sort to do the actual job.
func (sorter *Sorter) SortByCurrentColumn(stocks []Stock) *Sorter {
	var interfaces []sort.Interface

	if sorter.profile.Ascending {
		interfaces = []sort.Interface{
			byTickerAsc       { stocks },
			byLastTradeAsc    { stocks },
			byChangeAsc       { stocks },
			byChangePctAsc    { stocks },
			byOpenAsc         { stocks },
			byLowAsc          { stocks },
			byHighAsc         { stocks },
			byLow52Asc        { stocks },
			byHigh52Asc       { stocks },
			byVolumeAsc       { stocks },
			byAvgVolumeAsc    { stocks },
			byPeRatioAsc      { stocks },
			byDividendAsc     { stocks },
			byYieldAsc        { stocks },
			byMarketCapAsc    { stocks },
		}
	} else {
		interfaces = []sort.Interface{
			byTickerDesc      { stocks },
			byLastTradeDesc   { stocks },
			byChangeDesc      { stocks },
			byChangePctDesc   { stocks },
			byOpenDesc        { stocks },
			byLowDesc         { stocks },
			byHighDesc        { stocks },
			byLow52Desc       { stocks },
			byHigh52Desc      { stocks },
			byVolumeDesc      { stocks },
			byAvgVolumeDesc   { stocks },
			byPeRatioDesc     { stocks },
			byDividendDesc    { stocks },
			byYieldDesc       { stocks },
			byMarketCapDesc   { stocks },
		}
	}

	sort.Sort(interfaces[sorter.profile.SortColumn])

	return sorter
}

// The same exact method is used to sort by $Change and Change%. In both cases
// we sort by the value of Change% so that multiple $0.00s get sorted proferly.
func c(str string) float32 {
	trimmed := strings.Replace(strings.Trim(str, ` %`), `$`, ``, 1)
	value, _ := strconv.ParseFloat(trimmed, 32)
	return float32(value)
}

// When sorting by the market value we must first convert 42B etc. notations
// to proper numeric values.
func m(str string) float32 {
	multiplier := 1.0

	switch str[len(str)-1:len(str)] {	// Check the last character.
	case `B`:
		multiplier = 1000000000.0
	case `M`:
		multiplier = 1000000.0
	case `K`:
		multiplier = 1000.0
	}

	trimmed := strings.Trim(str, ` $BMK`)	// Get rid of non-numeric characters.
	value, _ := strconv.ParseFloat(trimmed, 32)

	return float32(value * multiplier)
}
<|endoftext|>


<commit_before>package sphere

import (
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"github.com/rs/xid"
)

const (
	// Read buffer size for websocket upgrader
	readBufferSize = 1024
	// Write buffer size for websocker upgrader
	writeBufferSize = 1024
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second
	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second
	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

var (
	// Guid to generate globally unique id
	guid = xid.New()
)

// Default creates a new instance of Sphere
func Default(opts ...interface{}) *Sphere {
	// declare agent
	var broker IBroker
	var option *Option
	// set declared agent if parameter exists
	for _, i := range opts {
		switch obj := i.(type) {
		case IBroker:
			broker = obj
		case *Option:
			option = obj
		}
	}
	if broker == nil {
		broker = DefaultSimpleBroker()
	}
	// websocket upgrader
	upgrader := websocket.Upgrader{ReadBufferSize: readBufferSize, WriteBufferSize: writeBufferSize}
	// update websocket upgrader with option object
	if option != nil && option.CheckOrigin {
		upgrader.CheckOrigin = func(r *http.Request) bool {
			return !option.CheckOrigin
		}
	}
	// creates sphere instance
	sphere := &Sphere{
		broker:      broker,
		connections: newConnectionMap(),
		channels:    newChannelMap(),
		models:      newChannelModelMap(),
		events:      newEventModelMap(),
		upgrader:    upgrader,
	}
	return sphere
}

// Sphere represents an entire Websocket instance
type Sphere struct {
	// a broker agent
	broker IBroker
	// list of active connections
	connections connectionmap
	// list of channels
	channels channelmap
	// list of models
	models channelmodelmap
	// list of events
	events eventmodelmap
	// websocket upgrader
	upgrader websocket.Upgrader
}

// Option for Sphere
type Option struct {
	CheckOrigin bool
}

// Handler handles and creates websocket connection
func (sphere *Sphere) Handler(w http.ResponseWriter, r *http.Request) IError {
	if conn, err := NewConnection(sphere.upgrader, w, r); err == nil {
		sphere.connections.Set(conn.id, conn)
		// run connection queue
		go conn.queue()
		// action after connection disconnected
		defer func() {
			// unsubscribe all channels
			for item := range conn.channels.Iter() {
				channel := item.Val
				sphere.unsubscribe(channel.namespace, channel.room, conn)
			}
			// close all send and receive buffers
			conn.close()
			// remove connection from sphere after disconnect
			sphere.connections.Remove(conn.id)
		}()
		for {
			_, msg, err := conn.ReadMessage()
			if err != nil {
				return err
			}
			if msg != nil {
				go sphere.process(conn, msg)
			}
		}
	} else {
		return err
	}
	return nil
}

// Models load channel or event models
func (sphere *Sphere) Models(models ...interface{}) {
	for _, item := range models {
		switch model := item.(type) {
		case IChannels:
			if !sphere.models.Has(model.Namespace()) {
				sphere.models.Set(model.Namespace(), model)
			}
		case IEvents:
			if !sphere.events.Has(model.Namespace()) {
				sphere.events.Set(model.Namespace(), model)
			}
		}
	}
}

// process parses and processes received message
func (sphere *Sphere) process(conn *Connection, msg []byte) {
	// convert received bytes to Packet object
	p, err := ParsePacket(msg)
	if err != nil {
		LogError(err)
		return
	}
	switch p.Type {
	case PacketTypeChannel:
		if p.Namespace != "" && p.Room != "" {
			// publish message to broker if it is a channel event / message
			p.Machine = sphere.broker.ID()
			sphere.publish(p, conn)
		} else {
			// if namespace or room is not provided, return error message
			p.Error = ErrBadScheme
			conn.send <- p
		}
	case PacketTypeSubscribe:
		if p.Namespace != "" && p.Room != "" {
			// subscribe connection to channel
			err := sphere.subscribe(p.Namespace, p.Room, conn)
			r := p.Response()
			r.SetError(err)
			// return success or failure message to user
			conn.send <- r
		} else {
			// if namespace or room is not provided, return error message
			p.Error = ErrBadScheme
			conn.send <- p
		}
	case PacketTypeUnsubscribe:
		if p.Namespace != "" && p.Room != "" {
			// unsubscribe connection from channel
			if sphere.models.Has(p.Namespace) {
				sphere.unsubscribe(p.Namespace, p.Room, conn)
			} else {
				// if namespace model does not existed, return error
				p.Error = ErrNotSupported
				conn.send <- p
			}
		} else {
			// if namespace or room is not provided, return error message
			p.Error = ErrBadScheme
			conn.send <- p
		}
	case PacketTypeMessage:
		if p.Namespace != "" {
			// receive event message
			sphere.receive(p, conn)
		} else {
			p.Error = ErrBadScheme
			conn.send <- p
		}
	case PacketTypePing:
		// ping-pong
		r := p.Response()
		conn.send <- r
	}
}

// channel returns Channel object, channel will be automatually created when autoCreateOpts is true
func (sphere *Sphere) channel(namespace string, room string, autoCreateOpts ...bool) *Channel {
	c := make(chan *Channel)
	autoCreateOpt := false
	for _, opt := range autoCreateOpts {
		autoCreateOpt = opt
		break
	}
	name := sphere.broker.ChannelName(namespace, room)
	go func() {
		if tmp, ok := sphere.channels.Get(name); ok {
			c <- tmp
		} else {
			if autoCreateOpt {
				channel := NewChannel(namespace, room)
				sphere.channels.Set(name, channel)
				c <- channel
			} else {
				c <- nil
			}
		}
	}()
	return <-c
}

// subscribe trigger Broker OnSubscribe action and put connection into channel connections list
func (sphere *Sphere) subscribe(namespace string, room string, conn *Connection) IError {
	var model IChannels
	if !sphere.models.Has(namespace) {
		return ErrNotSupported
	}
	if tmp, ok := sphere.models.Get(namespace); ok {
		model = tmp
	} else {
		return ErrNotSupported
	}
	if accept, err := model.Subscribe(room, conn); !accept && err == nil {
		return ErrUnauthorized
	} else if !accept && err != nil {
		return err
	}
	channel := sphere.channel(namespace, room, true)
	if channel == nil {
		return ErrNotFound
	}
	if err := channel.subscribe(conn); err != nil {
		return err
	}
	if !sphere.broker.IsSubscribed(channel.namespace, channel.room) {
		c := make(chan IError)
		go sphere.broker.OnSubscribe(channel, c)
		return <-c
	}
	return nil
}

// unsubscribe trigger Broker OnUnsubscribe action and remove connection from channel connections list
func (sphere *Sphere) unsubscribe(namespace string, room string, conn *Connection) IError {
	var model IChannels
	if !sphere.models.Has(namespace) {
		return ErrNotSupported
	}
	if tmp, ok := sphere.models.Get(namespace); ok {
		model = tmp
	} else {
		return ErrNotSupported
	}
	if err := model.Disconnect(room, conn); err != nil {
		return err
	}
	channel := sphere.channel(namespace, room, false)
	if channel == nil {
		return ErrNotFound
	}
	err := channel.unsubscribe(conn)
	if err == nil && channel.connections.Count() == 0 {
		if sphere.broker.IsSubscribed(channel.namespace, channel.room) {
			c := make(chan IError)
			go sphere.broker.OnUnsubscribe(channel, c)
			return <-c
		}
	}
	return err
}

// publish trigger Broker OnPublish action, send message to user from broker
func (sphere *Sphere) publish(p *Packet, conn *Connection) IError {
	var model IChannels
	if !sphere.models.Has(p.Namespace) {
		return ErrNotSupported
	}
	if tmp, ok := sphere.models.Get(p.Namespace); ok {
		model = tmp
	} else {
		return ErrNotSupported
	}
	channel := sphere.channel(p.Namespace, p.Room)
	if channel == nil {
		return ErrBadStatus
	}
	if !channel.isSubscribed(conn) {
		return ErrNotSubscribed
	}
	msg := p.Message
	if msg == nil || msg.Event == "" {
		return ErrBadScheme
	}
	res, err := model.Receive(msg.Event, msg.Data)
	if err != nil {
		return err
	}
	d := p.Response()
	if res != "" {
		d.Message.Data = res
	}
	if sphere.broker.IsSubscribed(channel.namespace, channel.room) {
		return sphere.broker.OnPublish(channel, d)
	}
	return ErrServerErrors
}

// receive message and event handler
func (sphere *Sphere) receive(p *Packet, conn *Connection) IError {
	var model IEvents
	if !sphere.events.Has(p.Namespace) {
		return ErrNotSupported
	}
	if tmp, ok := sphere.events.Get(p.Namespace); ok {
		model = tmp
	} else {
		return ErrNotSupported
	}
	msg := p.Message
	if msg == nil || msg.Event == "" {
		return ErrBadScheme
	}
	res, err := model.Receive(msg.Event, msg.Data)
	if err != nil {
		return err
	}
	d := p.Response()
	if res != "" {
		d.Message.Data = res
	}
	conn.send <- d
	return nil
}
<commit_msg>fix opts CheckOrigin "true" only, should accept false value<commit_after>package sphere

import (
	"net/http"
	"time"

	"github.com/gorilla/websocket"
	"github.com/rs/xid"
)

const (
	// Read buffer size for websocket upgrader
	readBufferSize = 1024
	// Write buffer size for websocker upgrader
	writeBufferSize = 1024
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second
	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second
	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10
	// Maximum message size allowed from peer.
	maxMessageSize = 512
)

var (
	// Guid to generate globally unique id
	guid = xid.New()
)

// Default creates a new instance of Sphere
func Default(opts ...interface{}) *Sphere {
	// declare agent
	var broker IBroker
	var option *Option
	// set declared agent if parameter exists
	for _, i := range opts {
		switch obj := i.(type) {
		case IBroker:
			broker = obj
		case *Option:
			option = obj
		}
	}
	if broker == nil {
		broker = DefaultSimpleBroker()
	}
	// websocket upgrader
	upgrader := websocket.Upgrader{ReadBufferSize: readBufferSize, WriteBufferSize: writeBufferSize}
	// update websocket upgrader with option object
	if option != nil {
		upgrader.CheckOrigin = func(r *http.Request) bool {
			return !option.CheckOrigin
		}
	}
	// creates sphere instance
	sphere := &Sphere{
		broker:      broker,
		connections: newConnectionMap(),
		channels:    newChannelMap(),
		models:      newChannelModelMap(),
		events:      newEventModelMap(),
		upgrader:    upgrader,
	}
	return sphere
}

// Sphere represents an entire Websocket instance
type Sphere struct {
	// a broker agent
	broker IBroker
	// list of active connections
	connections connectionmap
	// list of channels
	channels channelmap
	// list of models
	models channelmodelmap
	// list of events
	events eventmodelmap
	// websocket upgrader
	upgrader websocket.Upgrader
}

// Option for Sphere
type Option struct {
	CheckOrigin bool
}

// Handler handles and creates websocket connection
func (sphere *Sphere) Handler(w http.ResponseWriter, r *http.Request) IError {
	if conn, err := NewConnection(sphere.upgrader, w, r); err == nil {
		sphere.connections.Set(conn.id, conn)
		// run connection queue
		go conn.queue()
		// action after connection disconnected
		defer func() {
			// unsubscribe all channels
			for item := range conn.channels.Iter() {
				channel := item.Val
				sphere.unsubscribe(channel.namespace, channel.room, conn)
			}
			// close all send and receive buffers
			conn.close()
			// remove connection from sphere after disconnect
			sphere.connections.Remove(conn.id)
		}()
		for {
			_, msg, err := conn.ReadMessage()
			if err != nil {
				return err
			}
			if msg != nil {
				go sphere.process(conn, msg)
			}
		}
	} else {
		return err
	}
	return nil
}

// Models load channel or event models
func (sphere *Sphere) Models(models ...interface{}) {
	for _, item := range models {
		switch model := item.(type) {
		case IChannels:
			if !sphere.models.Has(model.Namespace()) {
				sphere.models.Set(model.Namespace(), model)
			}
		case IEvents:
			if !sphere.events.Has(model.Namespace()) {
				sphere.events.Set(model.Namespace(), model)
			}
		}
	}
}

// process parses and processes received message
func (sphere *Sphere) process(conn *Connection, msg []byte) {
	// convert received bytes to Packet object
	p, err := ParsePacket(msg)
	if err != nil {
		LogError(err)
		return
	}
	switch p.Type {
	case PacketTypeChannel:
		if p.Namespace != "" && p.Room != "" {
			// publish message to broker if it is a channel event / message
			p.Machine = sphere.broker.ID()
			sphere.publish(p, conn)
		} else {
			// if namespace or room is not provided, return error message
			p.Error = ErrBadScheme
			conn.send <- p
		}
	case PacketTypeSubscribe:
		if p.Namespace != "" && p.Room != "" {
			// subscribe connection to channel
			err := sphere.subscribe(p.Namespace, p.Room, conn)
			r := p.Response()
			r.SetError(err)
			// return success or failure message to user
			conn.send <- r
		} else {
			// if namespace or room is not provided, return error message
			p.Error = ErrBadScheme
			conn.send <- p
		}
	case PacketTypeUnsubscribe:
		if p.Namespace != "" && p.Room != "" {
			// unsubscribe connection from channel
			if sphere.models.Has(p.Namespace) {
				sphere.unsubscribe(p.Namespace, p.Room, conn)
			} else {
				// if namespace model does not existed, return error
				p.Error = ErrNotSupported
				conn.send <- p
			}
		} else {
			// if namespace or room is not provided, return error message
			p.Error = ErrBadScheme
			conn.send <- p
		}
	case PacketTypeMessage:
		if p.Namespace != "" {
			// receive event message
			sphere.receive(p, conn)
		} else {
			p.Error = ErrBadScheme
			conn.send <- p
		}
	case PacketTypePing:
		// ping-pong
		r := p.Response()
		conn.send <- r
	}
}

// channel returns Channel object, channel will be automatually created when autoCreateOpts is true
func (sphere *Sphere) channel(namespace string, room string, autoCreateOpts ...bool) *Channel {
	c := make(chan *Channel)
	autoCreateOpt := false
	for _, opt := range autoCreateOpts {
		autoCreateOpt = opt
		break
	}
	name := sphere.broker.ChannelName(namespace, room)
	go func() {
		if tmp, ok := sphere.channels.Get(name); ok {
			c <- tmp
		} else {
			if autoCreateOpt {
				channel := NewChannel(namespace, room)
				sphere.channels.Set(name, channel)
				c <- channel
			} else {
				c <- nil
			}
		}
	}()
	return <-c
}

// subscribe trigger Broker OnSubscribe action and put connection into channel connections list
func (sphere *Sphere) subscribe(namespace string, room string, conn *Connection) IError {
	var model IChannels
	if !sphere.models.Has(namespace) {
		return ErrNotSupported
	}
	if tmp, ok := sphere.models.Get(namespace); ok {
		model = tmp
	} else {
		return ErrNotSupported
	}
	if accept, err := model.Subscribe(room, conn); !accept && err == nil {
		return ErrUnauthorized
	} else if !accept && err != nil {
		return err
	}
	channel := sphere.channel(namespace, room, true)
	if channel == nil {
		return ErrNotFound
	}
	if err := channel.subscribe(conn); err != nil {
		return err
	}
	if !sphere.broker.IsSubscribed(channel.namespace, channel.room) {
		c := make(chan IError)
		go sphere.broker.OnSubscribe(channel, c)
		return <-c
	}
	return nil
}

// unsubscribe trigger Broker OnUnsubscribe action and remove connection from channel connections list
func (sphere *Sphere) unsubscribe(namespace string, room string, conn *Connection) IError {
	var model IChannels
	if !sphere.models.Has(namespace) {
		return ErrNotSupported
	}
	if tmp, ok := sphere.models.Get(namespace); ok {
		model = tmp
	} else {
		return ErrNotSupported
	}
	if err := model.Disconnect(room, conn); err != nil {
		return err
	}
	channel := sphere.channel(namespace, room, false)
	if channel == nil {
		return ErrNotFound
	}
	err := channel.unsubscribe(conn)
	if err == nil && channel.connections.Count() == 0 {
		if sphere.broker.IsSubscribed(channel.namespace, channel.room) {
			c := make(chan IError)
			go sphere.broker.OnUnsubscribe(channel, c)
			return <-c
		}
	}
	return err
}

// publish trigger Broker OnPublish action, send message to user from broker
func (sphere *Sphere) publish(p *Packet, conn *Connection) IError {
	var model IChannels
	if !sphere.models.Has(p.Namespace) {
		return ErrNotSupported
	}
	if tmp, ok := sphere.models.Get(p.Namespace); ok {
		model = tmp
	} else {
		return ErrNotSupported
	}
	channel := sphere.channel(p.Namespace, p.Room)
	if channel == nil {
		return ErrBadStatus
	}
	if !channel.isSubscribed(conn) {
		return ErrNotSubscribed
	}
	msg := p.Message
	if msg == nil || msg.Event == "" {
		return ErrBadScheme
	}
	res, err := model.Receive(msg.Event, msg.Data)
	if err != nil {
		return err
	}
	d := p.Response()
	if res != "" {
		d.Message.Data = res
	}
	if sphere.broker.IsSubscribed(channel.namespace, channel.room) {
		return sphere.broker.OnPublish(channel, d)
	}
	return ErrServerErrors
}

// receive message and event handler
func (sphere *Sphere) receive(p *Packet, conn *Connection) IError {
	var model IEvents
	if !sphere.events.Has(p.Namespace) {
		return ErrNotSupported
	}
	if tmp, ok := sphere.events.Get(p.Namespace); ok {
		model = tmp
	} else {
		return ErrNotSupported
	}
	msg := p.Message
	if msg == nil || msg.Event == "" {
		return ErrBadScheme
	}
	res, err := model.Receive(msg.Event, msg.Data)
	if err != nil {
		return err
	}
	d := p.Response()
	if res != "" {
		d.Message.Data = res
	}
	conn.send <- d
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"log"
	"time"

	"github.com/bigwhite/gocmpp"
)

const (
	userS     string = "900001"
	passwordS string = "888888"
)

func handleLogin(r *cmpp.Response, p *cmpp.Packet, l *log.Logger) (bool, error) {
	req, ok := p.Packer.(*cmpp.CmppConnReqPkt)
	if !ok {
		// not a connect request, ignore it,
		// go on to next handler
		return true, nil
	}

	resp := r.Packer.(*cmpp.Cmpp3ConnRspPkt)

	// validate the user and password
	// set the status in the connect response.
	resp.Version = 0x30
	addr := req.SrcAddr
	if addr != userS {
		l.Println("handleLogin error:", cmpp.ConnRspStatusErrMap[cmpp.ErrnoConnInvalidSrcAddr])
		resp.Status = uint32(cmpp.ErrnoConnInvalidSrcAddr)
		return false, cmpp.ConnRspStatusErrMap[cmpp.ErrnoConnInvalidSrcAddr]
	}

	tm := req.Timestamp
	authSrc := md5.Sum(bytes.Join([][]byte{[]byte(userS),
		make([]byte, 9),
		[]byte(passwordS),
		[]byte(fmt.Sprintf("%d", tm))},
		nil))

	if req.AuthSrc != string(authSrc[:]) {
		l.Println("handleLogin error: ", cmpp.ConnRspStatusErrMap[cmpp.ErrnoConnAuthFailed])
		resp.Status = uint32(cmpp.ErrnoConnAuthFailed)
		return false, cmpp.ConnRspStatusErrMap[cmpp.ErrnoConnAuthFailed]
	}

	authIsmg := md5.Sum(bytes.Join([][]byte{[]byte{byte(resp.Status)},
		authSrc[:],
		[]byte(passwordS)},
		nil))
	resp.AuthIsmg = string(authIsmg[:])
	l.Printf("handleLogin: %s login ok\n", addr)

	return false, nil
}

func handleSubmit(r *cmpp.Response, p *cmpp.Packet, l *log.Logger) (bool, error) {
	req, ok := p.Packer.(*cmpp.Cmpp3SubmitReqPkt)
	if !ok {
		return true, nil // go on to next handler
	}

	resp := r.Packer.(*cmpp.Cmpp3SubmitRspPkt)
	resp.MsgId = 12878564852733378560 //0xb2, 0xb9, 0xda, 0x80, 0x00, 0x01, 0x00, 0x00
	for i, d := range req.DestTerminalId {
		l.Printf("handleSubmit: handle submit from %s ok! msgid[%d], srcId[%s], destTerminalId[%s]\n",
			req.MsgSrc, resp.MsgId+uint64(i), req.SrcId, d)
	}
	return true, nil
}

func main() {
	var handlers = []cmpp.Handler{
		cmpp.HandlerFunc(handleLogin),
		cmpp.HandlerFunc(handleSubmit),
	}

	err := cmpp.ListenAndServe(":8888",
		cmpp.V30,
		5*time.Second,
		3,
		nil,
		handlers...,
	)
	if err != nil {
		log.Println("cmpp ListenAndServ error:", err)
	}
	return
}
<commit_msg>fix #2: use cmpputils.TimeStamp2Str to replace sprintf<commit_after>package main

import (
	"bytes"
	"crypto/md5"
	"log"
	"time"

	"github.com/bigwhite/gocmpp"
	"github.com/bigwhite/gocmpp/utils"
)

const (
	userS     string = "900001"
	passwordS string = "888888"
)

func handleLogin(r *cmpp.Response, p *cmpp.Packet, l *log.Logger) (bool, error) {
	req, ok := p.Packer.(*cmpp.CmppConnReqPkt)
	if !ok {
		// not a connect request, ignore it,
		// go on to next handler
		return true, nil
	}

	resp := r.Packer.(*cmpp.Cmpp3ConnRspPkt)

	// validate the user and password
	// set the status in the connect response.
	resp.Version = 0x30
	addr := req.SrcAddr
	if addr != userS {
		l.Println("handleLogin error:", cmpp.ConnRspStatusErrMap[cmpp.ErrnoConnInvalidSrcAddr])
		resp.Status = uint32(cmpp.ErrnoConnInvalidSrcAddr)
		return false, cmpp.ConnRspStatusErrMap[cmpp.ErrnoConnInvalidSrcAddr]
	}

	tm := req.Timestamp
	authSrc := md5.Sum(bytes.Join([][]byte{[]byte(userS),
		make([]byte, 9),
		[]byte(passwordS),
		[]byte(cmpputils.TimeStamp2Str(tm))},
		nil))

	if req.AuthSrc != string(authSrc[:]) {
		l.Println("handleLogin error: ", cmpp.ConnRspStatusErrMap[cmpp.ErrnoConnAuthFailed])
		resp.Status = uint32(cmpp.ErrnoConnAuthFailed)
		return false, cmpp.ConnRspStatusErrMap[cmpp.ErrnoConnAuthFailed]
	}

	authIsmg := md5.Sum(bytes.Join([][]byte{[]byte{byte(resp.Status)},
		authSrc[:],
		[]byte(passwordS)},
		nil))
	resp.AuthIsmg = string(authIsmg[:])
	l.Printf("handleLogin: %s login ok\n", addr)

	return false, nil
}

func handleSubmit(r *cmpp.Response, p *cmpp.Packet, l *log.Logger) (bool, error) {
	req, ok := p.Packer.(*cmpp.Cmpp3SubmitReqPkt)
	if !ok {
		return true, nil // go on to next handler
	}

	resp := r.Packer.(*cmpp.Cmpp3SubmitRspPkt)
	resp.MsgId = 12878564852733378560 //0xb2, 0xb9, 0xda, 0x80, 0x00, 0x01, 0x00, 0x00
	for i, d := range req.DestTerminalId {
		l.Printf("handleSubmit: handle submit from %s ok! msgid[%d], srcId[%s], destTerminalId[%s]\n",
			req.MsgSrc, resp.MsgId+uint64(i), req.SrcId, d)
	}
	return true, nil
}

func main() {
	var handlers = []cmpp.Handler{
		cmpp.HandlerFunc(handleLogin),
		cmpp.HandlerFunc(handleSubmit),
	}

	err := cmpp.ListenAndServe(":8888",
		cmpp.V30,
		5*time.Second,
		3,
		nil,
		handlers...,
	)
	if err != nil {
		log.Println("cmpp ListenAndServ error:", err)
	}
	return
}
<|endoftext|>


<commit_before>package db

// Level DB wrapper over multiple queue data store.
// Each key has a prefix that combines queue name as well as type of stored data.
// The following format is used for item keys:
// somename:m:itemid
// This format is used for payload keys:
// somename:p:payload

import (
	"bytes"
	"firempq/common"
	"firempq/qerrors"
	"firempq/util"
	"github.com/jmhodges/levigo"
	"github.com/op/go-logging"
	"sync"
	"time"
)

var log = logging.MustGetLogger("ldb")

// Item iterator built on top of LevelDB.
// It takes into account queue name to limit the amount of selected data.
type ItemIterator struct {
	iter   *levigo.Iterator
	prefix []byte // Prefix for look ups.
	Key    []byte // Currently selected key. Valid only if the iterator is valid.
	Value  []byte // Currently selected value. Valid only if the iterator is valid.
}

func NewIter(iter *levigo.Iterator, prefix []byte) *ItemIterator {
	iter.Seek(prefix)
	return &ItemIterator{iter, prefix, nil, nil}
}

// Switch to the next element.
func (mi *ItemIterator) Next() {
	mi.iter.Next()
}

// Before value is read, check if iterator is valid.
// for iter.Valid() {
//    mykey := iter.Key
//    myvalue := iter.Value
//    ......
//    iter.Next()
//}
func (mi *ItemIterator) Valid() bool {
	valid := mi.iter.Valid()
	if valid {
		k := mi.iter.Key()
		// Strip key prefix. If prefix doesn't match the length of the slice will remain the same.
		normKey := bytes.TrimPrefix(k, mi.prefix)
		if len(normKey) == len(k) {
			return false
		}
		mi.Key = normKey
		mi.Value = mi.iter.Value()
		return true
	}
	return false
}

// Iterator must be closed!
func (mi *ItemIterator) Close() {
	mi.iter.Close()
}

type DataStorage struct {
	db              *levigo.DB        // Pointer the the instance of level db.
	dbName          string            // LevelDB database name.
	itemCache       map[string][]byte // Active cache for item metadata.
	payloadCache    map[string]string // Active cache for item payload.
	tmpItemCache    map[string][]byte // Temporary map of cached data while it is in process of flushing into DB.
	tmpPayloadCache map[string]string // Temporary map of cached payloads while it is in process of flushing into DB.
	cacheLock       sync.Mutex        // Used for caches access.
	flushLock       sync.Mutex        // Used to prevent double flush.
	closed          bool
}

func NewDataStorage(dbName string) *DataStorage {
	ds := DataStorage{
		dbName:          dbName,
		itemCache:       make(map[string][]byte),
		payloadCache:    make(map[string]string),
		tmpItemCache:    make(map[string][]byte),
		tmpPayloadCache: make(map[string]string),
		closed:          false,
	}
	opts := levigo.NewOptions()
	opts.SetCreateIfMissing(true)
	opts.SetWriteBufferSize(10 * 1024 * 1024)
	opts.SetCompression(levigo.SnappyCompression)
	db, err := levigo.Open(dbName, opts)
	if err != nil {
		log.Critical("Could not initialize database: %s", err.Error())
		return nil
	}
	ds.db = db
	go ds.periodicCacheFlush()
	return &ds
}

func (ds *DataStorage) periodicCacheFlush() {
	for !ds.closed {
		ds.flushLock.Lock()
		if !ds.closed {
			ds.FlushCache()
		}
		ds.flushLock.Unlock()
		time.Sleep(100 * time.Millisecond)
	}
}

// Item payload Id.
func makePayloadId(queueName, id string) string {
	return "p:" + queueName + ":" + id
}

// Item Id.
func makeItemId(queueName, id string) string {
	return "m:" + queueName + ":" + id
}

// Item will be stored into cache including payload.
func (ds *DataStorage) StoreItemWithPayload(queueName string, item common.IItemMetaData, payload string) {
	itemId := makeItemId(queueName, item.GetId())
	payloadId := makePayloadId(queueName, item.GetId())

	itemBody := item.ToBinary()

	ds.cacheLock.Lock()
	ds.itemCache[itemId] = itemBody
	ds.payloadCache[payloadId] = payload
	ds.cacheLock.Unlock()
}

func (ds *DataStorage) StoreItem(queueName string, item common.IItemMetaData) {
	itemId := makeItemId(queueName, item.GetId())

	itemBody := item.ToBinary()

	ds.cacheLock.Lock()
	ds.itemCache[itemId] = itemBody
	ds.cacheLock.Unlock()
}

// Updates item metadata, affects cache only until flushed.
func (ds *DataStorage) UpdateItem(queueName string, item common.IItemMetaData) {
	itemId := makeItemId(queueName, item.GetId())
	itemBody := item.ToBinary()
	ds.cacheLock.Lock()
	ds.itemCache[itemId] = itemBody
	ds.cacheLock.Unlock()
}

// Deletes item metadata and payload, affects cache only until flushed.
func (ds *DataStorage) DeleteItem(queueName string, itemId string) {
	id := makeItemId(queueName, itemId)
	payloadId := makePayloadId(queueName, itemId)
	ds.cacheLock.Lock()
	ds.itemCache[id] = nil
	ds.payloadCache[payloadId] = ""
	ds.cacheLock.Unlock()
}

// Returns item payload. Three places are checked:
// 1. Top level cache.
// 2. Temp cache while data is getting flushed into db.
// 3. If not found in cache, will do actually DB lookup.
func (ds *DataStorage) GetPayload(queueName string, itemId string) string {
	payloadId := makePayloadId(queueName, itemId)
	ds.cacheLock.Lock()
	payload, ok := ds.payloadCache[itemId]
	if ok {
		ds.cacheLock.Unlock()
		return payload
	}
	payload, ok = ds.tmpPayloadCache[payloadId]
	if ok {
		ds.cacheLock.Unlock()
		return payload
	}
	ds.cacheLock.Unlock()
	ropts := levigo.NewReadOptions()
	value, _ := ds.db.Get(ropts, []byte(payloadId))
	return string(value)
}

func (ds *DataStorage) FlushCache() {
	if ds.closed {
		return
	}
	ds.cacheLock.Lock()
	ds.tmpItemCache = ds.itemCache
	ds.tmpPayloadCache = ds.payloadCache
	ds.itemCache = make(map[string][]byte)
	ds.payloadCache = make(map[string]string)
	ds.cacheLock.Unlock()

	wb := levigo.NewWriteBatch()
	for k, v := range ds.tmpItemCache {
		key := []byte(k)
		if v == nil {
			wb.Delete(key)
		} else {
			wb.Put(key, v)
		}
	}

	for k, v := range ds.tmpPayloadCache {
		key := []byte(k)
		if v == "" {
			wb.Delete(key)
		} else {
			wb.Put(key, []byte(v))
		}
	}

	wopts := levigo.NewWriteOptions()
	ds.db.Write(wopts, wb)

	ds.cacheLock.Lock()
	ds.tmpItemCache = make(map[string][]byte)
	ds.tmpPayloadCache = make(map[string]string)
	ds.cacheLock.Unlock()
}

// Returns new Iterator that must be closed!
// Queue name used as a prefix to iterate over all item metadata that belongs to the specific queue.
func (ds *DataStorage) IterQueue(queueName string) *ItemIterator {
	ropts := levigo.NewReadOptions()
	iter := ds.db.NewIterator(ropts)
	prefix := []byte(makeItemId(queueName, ""))
	return NewIter(iter, prefix)
}

const QUEUE_META_PREFIX = "qmeta:"
const QUEUE_SETTINGS_PREFIX = "qsettings:"

// Read all available queues.
func (ds *DataStorage) GetAllQueueMeta() []*common.QueueMetaInfo {
	qmiList := make([]*common.QueueMetaInfo, 0, 1000)
	qtPrefix := []byte(QUEUE_META_PREFIX)
	ropts := levigo.NewReadOptions()
	iter := ds.db.NewIterator(ropts)
	iter.Seek(qtPrefix)
	for iter.Valid() {
		queueKey := iter.Key()
		queueName := bytes.TrimPrefix(queueKey, qtPrefix)

		if len(queueKey) == len(queueName) {
			break
		}
		qmi, err := common.QueueInfoFromBinary(iter.Value())
		if err != nil {
			log.Error("Coudn't read queue meta data because of: %s", err.Error())
			iter.Next()
			continue
		}
		if qmi.Disabled {
			log.Error("Qeueue is disabled. Skipping: %s", qmi.Name)
			iter.Next()
			continue
		}
		qmiList = append(qmiList, qmi)
		iter.Next()
	}
	return qmiList
}

func (ds *DataStorage) SaveQueueMeta(qmi *common.QueueMetaInfo) {
	key := QUEUE_META_PREFIX + qmi.Name
	wopts := levigo.NewWriteOptions()
	ds.db.Put(wopts, []byte(key), qmi.ToBinary())
}

func makeSettingsKey(queueName string) []byte {
	return []byte(QUEUE_SETTINGS_PREFIX + queueName)
}

// Read queue settings bases on queue name. Caller should profide correct settings structure to read binary data.
func (ds *DataStorage) GetQueueSettings(settings interface{}, queueName string) error {
	key := makeSettingsKey(queueName)
	ropts := levigo.NewReadOptions()
	data, _ := ds.db.Get(ropts, key)
	if data == nil {
		return qerrors.InvalidRequest("No queue settings found: " + queueName)
	}
	err := util.StructFromBinary(settings, data)
	return err
}

func (ds *DataStorage) SaveQueueSettings(queueName string, settings interface{}) {
	key := makeSettingsKey(queueName)
	wopts := levigo.NewWriteOptions()
	data := util.StructToBinary(settings)
	err := ds.db.Put(wopts, key, data)
	if err != nil {
		log.Error("Failed to save settings: %s", err.Error())
	}
}

// Flush and close database.
func (ds *DataStorage) Close() {
	ds.flushLock.Lock()
	defer ds.flushLock.Unlock()
	if !ds.closed {
		ds.closed = true
		ds.FlushCache()
		ds.db.Close()
	} else {
		log.Error("Attempt to close database more than once!")
	}
}
<commit_msg>Added range compaction upon exits for level db. It may slow down the exit process. It might be better to run it on start.<commit_after>package db

// Level DB wrapper over multiple queue data store.
// Each key has a prefix that combines queue name as well as type of stored data.
// The following format is used for item keys:
// somename:m:itemid
// This format is used for payload keys:
// somename:p:payload

import (
	"bytes"
	"firempq/common"
	"firempq/qerrors"
	"firempq/util"
	"github.com/jmhodges/levigo"
	"github.com/op/go-logging"
	"sync"
	"time"
)

var log = logging.MustGetLogger("ldb")

// Item iterator built on top of LevelDB.
// It takes into account queue name to limit the amount of selected data.
type ItemIterator struct {
	iter   *levigo.Iterator
	prefix []byte // Prefix for look ups.
	Key    []byte // Currently selected key. Valid only if the iterator is valid.
	Value  []byte // Currently selected value. Valid only if the iterator is valid.
}

func NewIter(iter *levigo.Iterator, prefix []byte) *ItemIterator {
	iter.Seek(prefix)
	return &ItemIterator{iter, prefix, nil, nil}
}

// Switch to the next element.
func (mi *ItemIterator) Next() {
	mi.iter.Next()
}

// Before value is read, check if iterator is valid.
// for iter.Valid() {
//    mykey := iter.Key
//    myvalue := iter.Value
//    ......
//    iter.Next()
//}
func (mi *ItemIterator) Valid() bool {
	valid := mi.iter.Valid()
	if valid {
		k := mi.iter.Key()
		// Strip key prefix. If prefix doesn't match the length of the slice will remain the same.
		normKey := bytes.TrimPrefix(k, mi.prefix)
		if len(normKey) == len(k) {
			return false
		}
		mi.Key = normKey
		mi.Value = mi.iter.Value()
		return true
	}
	return false
}

// Iterator must be closed!
func (mi *ItemIterator) Close() {
	mi.iter.Close()
}

type DataStorage struct {
	db              *levigo.DB        // Pointer the the instance of level db.
	dbName          string            // LevelDB database name.
	itemCache       map[string][]byte // Active cache for item metadata.
	payloadCache    map[string]string // Active cache for item payload.
	tmpItemCache    map[string][]byte // Temporary map of cached data while it is in process of flushing into DB.
	tmpPayloadCache map[string]string // Temporary map of cached payloads while it is in process of flushing into DB.
	cacheLock       sync.Mutex        // Used for caches access.
	flushLock       sync.Mutex        // Used to prevent double flush.
	closed          bool
}

func NewDataStorage(dbName string) *DataStorage {
	ds := DataStorage{
		dbName:          dbName,
		itemCache:       make(map[string][]byte),
		payloadCache:    make(map[string]string),
		tmpItemCache:    make(map[string][]byte),
		tmpPayloadCache: make(map[string]string),
		closed:          false,
	}
	opts := levigo.NewOptions()
	opts.SetCreateIfMissing(true)
	opts.SetWriteBufferSize(10 * 1024 * 1024)
	opts.SetCompression(levigo.SnappyCompression)
	db, err := levigo.Open(dbName, opts)
	if err != nil {
		log.Critical("Could not initialize database: %s", err.Error())
		return nil
	}
	ds.db = db
	go ds.periodicCacheFlush()
	return &ds
}

func (ds *DataStorage) periodicCacheFlush() {
	for !ds.closed {
		ds.flushLock.Lock()
		if !ds.closed {
			ds.FlushCache()
		}
		ds.flushLock.Unlock()
		time.Sleep(100 * time.Millisecond)
	}
}

// Item payload Id.
func makePayloadId(queueName, id string) string {
	return "p:" + queueName + ":" + id
}

// Item Id.
func makeItemId(queueName, id string) string {
	return "m:" + queueName + ":" + id
}

// Item will be stored into cache including payload.
func (ds *DataStorage) StoreItemWithPayload(queueName string, item common.IItemMetaData, payload string) {
	itemId := makeItemId(queueName, item.GetId())
	payloadId := makePayloadId(queueName, item.GetId())

	itemBody := item.ToBinary()

	ds.cacheLock.Lock()
	ds.itemCache[itemId] = itemBody
	ds.payloadCache[payloadId] = payload
	ds.cacheLock.Unlock()
}

func (ds *DataStorage) StoreItem(queueName string, item common.IItemMetaData) {
	itemId := makeItemId(queueName, item.GetId())

	itemBody := item.ToBinary()

	ds.cacheLock.Lock()
	ds.itemCache[itemId] = itemBody
	ds.cacheLock.Unlock()
}

// Updates item metadata, affects cache only until flushed.
func (ds *DataStorage) UpdateItem(queueName string, item common.IItemMetaData) {
	itemId := makeItemId(queueName, item.GetId())
	itemBody := item.ToBinary()
	ds.cacheLock.Lock()
	ds.itemCache[itemId] = itemBody
	ds.cacheLock.Unlock()
}

// Deletes item metadata and payload, affects cache only until flushed.
func (ds *DataStorage) DeleteItem(queueName string, itemId string) {
	id := makeItemId(queueName, itemId)
	payloadId := makePayloadId(queueName, itemId)
	ds.cacheLock.Lock()
	ds.itemCache[id] = nil
	ds.payloadCache[payloadId] = ""
	ds.cacheLock.Unlock()
}

// Returns item payload. Three places are checked:
// 1. Top level cache.
// 2. Temp cache while data is getting flushed into db.
// 3. If not found in cache, will do actually DB lookup.
func (ds *DataStorage) GetPayload(queueName string, itemId string) string {
	payloadId := makePayloadId(queueName, itemId)
	ds.cacheLock.Lock()
	payload, ok := ds.payloadCache[itemId]
	if ok {
		ds.cacheLock.Unlock()
		return payload
	}
	payload, ok = ds.tmpPayloadCache[payloadId]
	if ok {
		ds.cacheLock.Unlock()
		return payload
	}
	ds.cacheLock.Unlock()
	ropts := levigo.NewReadOptions()
	value, _ := ds.db.Get(ropts, []byte(payloadId))
	return string(value)
}

func (ds *DataStorage) FlushCache() {
	if ds.closed {
		return
	}
	ds.cacheLock.Lock()
	ds.tmpItemCache = ds.itemCache
	ds.tmpPayloadCache = ds.payloadCache
	ds.itemCache = make(map[string][]byte)
	ds.payloadCache = make(map[string]string)
	ds.cacheLock.Unlock()

	wb := levigo.NewWriteBatch()
	for k, v := range ds.tmpItemCache {
		key := []byte(k)
		if v == nil {
			wb.Delete(key)
		} else {
			wb.Put(key, v)
		}
	}

	for k, v := range ds.tmpPayloadCache {
		key := []byte(k)
		if v == "" {
			wb.Delete(key)
		} else {
			wb.Put(key, []byte(v))
		}
	}

	wopts := levigo.NewWriteOptions()
	ds.db.Write(wopts, wb)

	ds.cacheLock.Lock()
	ds.tmpItemCache = make(map[string][]byte)
	ds.tmpPayloadCache = make(map[string]string)
	ds.cacheLock.Unlock()
}

// Returns new Iterator that must be closed!
// Queue name used as a prefix to iterate over all item metadata that belongs to the specific queue.
func (ds *DataStorage) IterQueue(queueName string) *ItemIterator {
	ropts := levigo.NewReadOptions()
	iter := ds.db.NewIterator(ropts)
	prefix := []byte(makeItemId(queueName, ""))
	return NewIter(iter, prefix)
}

const QUEUE_META_PREFIX = "qmeta:"
const QUEUE_SETTINGS_PREFIX = "qsettings:"

// Read all available queues.
func (ds *DataStorage) GetAllQueueMeta() []*common.QueueMetaInfo {
	qmiList := make([]*common.QueueMetaInfo, 0, 1000)
	qtPrefix := []byte(QUEUE_META_PREFIX)
	ropts := levigo.NewReadOptions()
	iter := ds.db.NewIterator(ropts)
	iter.Seek(qtPrefix)
	for iter.Valid() {
		queueKey := iter.Key()
		queueName := bytes.TrimPrefix(queueKey, qtPrefix)

		if len(queueKey) == len(queueName) {
			break
		}
		qmi, err := common.QueueInfoFromBinary(iter.Value())
		if err != nil {
			log.Error("Coudn't read queue meta data because of: %s", err.Error())
			iter.Next()
			continue
		}
		if qmi.Disabled {
			log.Error("Qeueue is disabled. Skipping: %s", qmi.Name)
			iter.Next()
			continue
		}
		qmiList = append(qmiList, qmi)
		iter.Next()
	}
	return qmiList
}

func (ds *DataStorage) SaveQueueMeta(qmi *common.QueueMetaInfo) {
	key := QUEUE_META_PREFIX + qmi.Name
	wopts := levigo.NewWriteOptions()
	ds.db.Put(wopts, []byte(key), qmi.ToBinary())
}

func makeSettingsKey(queueName string) []byte {
	return []byte(QUEUE_SETTINGS_PREFIX + queueName)
}

// Read queue settings bases on queue name. Caller should profide correct settings structure to read binary data.
func (ds *DataStorage) GetQueueSettings(settings interface{}, queueName string) error {
	key := makeSettingsKey(queueName)
	ropts := levigo.NewReadOptions()
	data, _ := ds.db.Get(ropts, key)
	if data == nil {
		return qerrors.InvalidRequest("No queue settings found: " + queueName)
	}
	err := util.StructFromBinary(settings, data)
	return err
}

func (ds *DataStorage) SaveQueueSettings(queueName string, settings interface{}) {
	key := makeSettingsKey(queueName)
	wopts := levigo.NewWriteOptions()
	data := util.StructToBinary(settings)
	err := ds.db.Put(wopts, key, data)
	if err != nil {
		log.Error("Failed to save settings: %s", err.Error())
	}
}

// Flush and close database.
func (ds *DataStorage) Close() {
	ds.flushLock.Lock()
	defer ds.flushLock.Unlock()
	if !ds.closed {
		ds.closed = true
		ds.FlushCache()
		ds.db.CompactRange(levigo.Range{nil, nil})
		ds.db.Close()
	} else {
		log.Error("Attempt to close database more than once!")
	}
}
<|endoftext|>


<commit_before>package scheduling

import (
	"bytes"
	"encoding/base64"
	"encoding/gob"
	"fmt"
	"path"
	"strings"
	"time"

	swarming_api "github.com/luci/luci-go/common/api/swarming/swarming/v1"
	"go.skia.org/infra/go/isolate"
	"go.skia.org/infra/go/swarming"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/task_scheduler/go/db"
	"go.skia.org/infra/task_scheduler/go/specs"
)

// taskCandidate is a struct used for determining which tasks to schedule.
type taskCandidate struct {
	Commits        []string
	IsolatedInput  string
	IsolatedHashes []string
	JobCreated     time.Time
	ParentTaskIds  []string
	RetryOf        string
	Score          float64
	StealingFromId string
	db.TaskKey
	TaskSpec *specs.TaskSpec
}

// Copy returns a copy of the taskCandidate.
func (c *taskCandidate) Copy() *taskCandidate {
	return &taskCandidate{
		Commits:        util.CopyStringSlice(c.Commits),
		IsolatedInput:  c.IsolatedInput,
		IsolatedHashes: util.CopyStringSlice(c.IsolatedHashes),
		JobCreated:     c.JobCreated,
		ParentTaskIds:  util.CopyStringSlice(c.ParentTaskIds),
		RetryOf:        c.RetryOf,
		Score:          c.Score,
		StealingFromId: c.StealingFromId,
		TaskKey:        c.TaskKey.Copy(),
		TaskSpec:       c.TaskSpec.Copy(),
	}
}

// MakeId generates a string ID for the taskCandidate.
func (c *taskCandidate) MakeId() string {
	var buf bytes.Buffer
	if err := gob.NewEncoder(&buf).Encode(&c.TaskKey); err != nil {
		panic(fmt.Sprintf("Failed to GOB encode TaskKey: %s", err))
	}
	b64 := base64.StdEncoding.EncodeToString(buf.Bytes())
	return fmt.Sprintf("taskCandidate|%s", b64)
}

// parseId generates taskCandidate information from the ID.
func parseId(id string) (db.TaskKey, error) {
	var rv db.TaskKey
	split := strings.Split(id, "|")
	if len(split) != 2 {
		return rv, fmt.Errorf("Invalid ID, not enough parts: %q", id)
	}
	if split[0] != "taskCandidate" {
		return rv, fmt.Errorf("Invalid ID, no 'taskCandidate' prefix: %q", id)
	}
	b, err := base64.StdEncoding.DecodeString(split[1])
	if err != nil {
		return rv, fmt.Errorf("Failed to base64 decode ID: %s", err)
	}
	if err := gob.NewDecoder(bytes.NewBuffer(b)).Decode(&rv); err != nil {
		return rv, fmt.Errorf("Failed to GOB decode ID: %s", err)
	}
	return rv, nil
}

// MakeTask instantiates a db.Task from the taskCandidate.
func (c *taskCandidate) MakeTask() *db.Task {
	commits := make([]string, len(c.Commits))
	copy(commits, c.Commits)
	parentTaskIds := make([]string, len(c.ParentTaskIds))
	copy(parentTaskIds, c.ParentTaskIds)
	return &db.Task{
		Commits:       commits,
		Id:            "", // Filled in when the task is inserted into the DB.
		ParentTaskIds: parentTaskIds,
		RetryOf:       c.RetryOf,
		TaskKey:       c.TaskKey.Copy(),
	}
}

// MakeIsolateTask creates an isolate.Task from this taskCandidate.
func (c *taskCandidate) MakeIsolateTask(infraBotsDir, baseDir string) *isolate.Task {
	return &isolate.Task{
		BaseDir:     baseDir,
		Blacklist:   isolate.DEFAULT_BLACKLIST,
		Deps:        c.IsolatedHashes,
		IsolateFile: path.Join(infraBotsDir, c.TaskSpec.Isolate),
		OsType:      "linux", // TODO(borenet)
	}
}

// getPatchStorage returns "gerrit" or "rietveld" based on the Server URL.
func getPatchStorage(server string) string {
	if server == "" {
		return ""
	}
	if strings.Contains(server, "codereview.chromium") {
		return "rietveld"
	}
	return "gerrit"
}

// replaceVars replaces variable names with their values in a given string.
func replaceVars(c *taskCandidate, s string) string {
	issueShort := ""
	if len(c.Issue) < specs.ISSUE_SHORT_LENGTH {
		issueShort = c.Issue
	} else {
		issueShort = c.Issue[len(c.Issue)-specs.ISSUE_SHORT_LENGTH:]
	}
	replacements := map[string]string{
		specs.VARIABLE_CODEREVIEW_SERVER: c.Server,
		specs.VARIABLE_ISSUE:             c.Issue,
		specs.VARIABLE_ISSUE_SHORT:       issueShort,
		specs.VARIABLE_PATCH_STORAGE:     getPatchStorage(c.Server),
		specs.VARIABLE_PATCHSET:          c.Patchset,
		specs.VARIABLE_REPO:              c.Repo,
		specs.VARIABLE_REVISION:          c.Revision,
		specs.VARIABLE_TASK_NAME:         c.Name,
	}
	for k, v := range replacements {
		s = strings.Replace(s, fmt.Sprintf(specs.VARIABLE_SYNTAX, k), v, -1)
	}
	return s
}

// MakeTaskRequest creates a SwarmingRpcsNewTaskRequest object from the taskCandidate.
func (c *taskCandidate) MakeTaskRequest(id string) *swarming_api.SwarmingRpcsNewTaskRequest {
	var cipdInput *swarming_api.SwarmingRpcsCipdInput
	if len(c.TaskSpec.CipdPackages) > 0 {
		cipdInput = &swarming_api.SwarmingRpcsCipdInput{
			Packages: make([]*swarming_api.SwarmingRpcsCipdPackage, 0, len(c.TaskSpec.CipdPackages)),
		}
		for _, p := range c.TaskSpec.CipdPackages {
			cipdInput.Packages = append(cipdInput.Packages, &swarming_api.SwarmingRpcsCipdPackage{
				PackageName: p.Name,
				Path:        p.Path,
				Version:     p.Version,
			})
		}
	}

	dims := make([]*swarming_api.SwarmingRpcsStringPair, 0, len(c.TaskSpec.Dimensions))
	dimsMap := make(map[string]string, len(c.TaskSpec.Dimensions))
	for _, d := range c.TaskSpec.Dimensions {
		split := strings.SplitN(d, ":", 2)
		key := split[0]
		val := split[1]
		dims = append(dims, &swarming_api.SwarmingRpcsStringPair{
			Key:   key,
			Value: val,
		})
		dimsMap[key] = val
	}

	var env []*swarming_api.SwarmingRpcsStringPair
	if len(c.TaskSpec.Environment) > 0 {
		env = make([]*swarming_api.SwarmingRpcsStringPair, 0, len(c.TaskSpec.Environment))
		for k, v := range c.TaskSpec.Environment {
			env = append(env, &swarming_api.SwarmingRpcsStringPair{
				Key:   k,
				Value: v,
			})
		}
	}

	extraArgs := make([]string, 0, len(c.TaskSpec.ExtraArgs))
	for _, arg := range c.TaskSpec.ExtraArgs {
		extraArgs = append(extraArgs, replaceVars(c, arg))
	}

	expirationSecs := int64(c.TaskSpec.Expiration.Seconds())
	if expirationSecs == int64(0) {
		expirationSecs = int64(swarming.RECOMMENDED_EXPIRATION.Seconds())
	}
	executionTimeoutSecs := int64(c.TaskSpec.ExecutionTimeout.Seconds())
	if executionTimeoutSecs == int64(0) {
		executionTimeoutSecs = int64(swarming.RECOMMENDED_HARD_TIMEOUT.Seconds())
	}
	ioTimeoutSecs := int64(c.TaskSpec.IoTimeout.Seconds())
	if ioTimeoutSecs == int64(0) {
		ioTimeoutSecs = int64(swarming.RECOMMENDED_IO_TIMEOUT.Seconds())
	}
	return &swarming_api.SwarmingRpcsNewTaskRequest{
		ExpirationSecs: expirationSecs,
		Name:           c.Name,
		Priority:       int64(100.0 * c.TaskSpec.Priority),
		Properties: &swarming_api.SwarmingRpcsTaskProperties{
			CipdInput:            cipdInput,
			Dimensions:           dims,
			Env:                  env,
			ExecutionTimeoutSecs: executionTimeoutSecs,
			ExtraArgs:            extraArgs,
			InputsRef: &swarming_api.SwarmingRpcsFilesRef{
				Isolated:       c.IsolatedInput,
				Isolatedserver: isolate.ISOLATE_SERVER_URL,
				Namespace:      isolate.DEFAULT_NAMESPACE,
			},
			IoTimeoutSecs: ioTimeoutSecs,
		},
		Tags: db.TagsForTask(c.Name, id, c.TaskSpec.Priority, c.RepoState, c.RetryOf, dimsMap, c.ForcedJobId, c.ParentTaskIds),
		User: "skia-task-scheduler",
	}
}

// allDepsMet determines whether all dependencies for the given task candidate
// have been satisfied, and if so, returns a map of whose keys are task IDs and
// values are their isolated outputs.
func (c *taskCandidate) allDepsMet(cache db.TaskCache) (bool, map[string]string, error) {
	rv := make(map[string]string, len(c.TaskSpec.Dependencies))
	for _, depName := range c.TaskSpec.Dependencies {
		key := c.TaskKey.Copy()
		key.Name = depName
		byKey, err := cache.GetTasksByKey(&key)
		if err != nil {
			return false, nil, err
		}
		ok := false
		for _, t := range byKey {
			if t.Done() && t.Success() && t.IsolatedOutput != "" {
				rv[t.Id] = t.IsolatedOutput
				ok = true
				break
			}
		}
		if !ok {
			return false, nil, nil
		}
	}
	return true, rv, nil
}

// taskCandidateSlice is an alias used for sorting a slice of taskCandidates.
type taskCandidateSlice []*taskCandidate

func (s taskCandidateSlice) Len() int { return len(s) }
func (s taskCandidateSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s taskCandidateSlice) Less(i, j int) bool {
	return s[i].Score > s[j].Score // candidates sort in decreasing order.
}
<commit_msg>[task scheduler] Get the isolate OsType from dimensions<commit_after>package scheduling

import (
	"bytes"
	"encoding/base64"
	"encoding/gob"
	"fmt"
	"path"
	"strings"
	"time"

	swarming_api "github.com/luci/luci-go/common/api/swarming/swarming/v1"
	"go.skia.org/infra/go/isolate"
	"go.skia.org/infra/go/swarming"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/task_scheduler/go/db"
	"go.skia.org/infra/task_scheduler/go/specs"
)

// taskCandidate is a struct used for determining which tasks to schedule.
type taskCandidate struct {
	Commits        []string
	IsolatedInput  string
	IsolatedHashes []string
	JobCreated     time.Time
	ParentTaskIds  []string
	RetryOf        string
	Score          float64
	StealingFromId string
	db.TaskKey
	TaskSpec *specs.TaskSpec
}

// Copy returns a copy of the taskCandidate.
func (c *taskCandidate) Copy() *taskCandidate {
	return &taskCandidate{
		Commits:        util.CopyStringSlice(c.Commits),
		IsolatedInput:  c.IsolatedInput,
		IsolatedHashes: util.CopyStringSlice(c.IsolatedHashes),
		JobCreated:     c.JobCreated,
		ParentTaskIds:  util.CopyStringSlice(c.ParentTaskIds),
		RetryOf:        c.RetryOf,
		Score:          c.Score,
		StealingFromId: c.StealingFromId,
		TaskKey:        c.TaskKey.Copy(),
		TaskSpec:       c.TaskSpec.Copy(),
	}
}

// MakeId generates a string ID for the taskCandidate.
func (c *taskCandidate) MakeId() string {
	var buf bytes.Buffer
	if err := gob.NewEncoder(&buf).Encode(&c.TaskKey); err != nil {
		panic(fmt.Sprintf("Failed to GOB encode TaskKey: %s", err))
	}
	b64 := base64.StdEncoding.EncodeToString(buf.Bytes())
	return fmt.Sprintf("taskCandidate|%s", b64)
}

// parseId generates taskCandidate information from the ID.
func parseId(id string) (db.TaskKey, error) {
	var rv db.TaskKey
	split := strings.Split(id, "|")
	if len(split) != 2 {
		return rv, fmt.Errorf("Invalid ID, not enough parts: %q", id)
	}
	if split[0] != "taskCandidate" {
		return rv, fmt.Errorf("Invalid ID, no 'taskCandidate' prefix: %q", id)
	}
	b, err := base64.StdEncoding.DecodeString(split[1])
	if err != nil {
		return rv, fmt.Errorf("Failed to base64 decode ID: %s", err)
	}
	if err := gob.NewDecoder(bytes.NewBuffer(b)).Decode(&rv); err != nil {
		return rv, fmt.Errorf("Failed to GOB decode ID: %s", err)
	}
	return rv, nil
}

// MakeTask instantiates a db.Task from the taskCandidate.
func (c *taskCandidate) MakeTask() *db.Task {
	commits := make([]string, len(c.Commits))
	copy(commits, c.Commits)
	parentTaskIds := make([]string, len(c.ParentTaskIds))
	copy(parentTaskIds, c.ParentTaskIds)
	return &db.Task{
		Commits:       commits,
		Id:            "", // Filled in when the task is inserted into the DB.
		ParentTaskIds: parentTaskIds,
		RetryOf:       c.RetryOf,
		TaskKey:       c.TaskKey.Copy(),
	}
}

// MakeIsolateTask creates an isolate.Task from this taskCandidate.
func (c *taskCandidate) MakeIsolateTask(infraBotsDir, baseDir string) *isolate.Task {
	os := "linux"
	for _, d := range c.TaskSpec.Dimensions {
		if strings.HasPrefix(d, "os:") {
			os = d[len("os:"):]
			break
		}
	}
	return &isolate.Task{
		BaseDir:     baseDir,
		Blacklist:   isolate.DEFAULT_BLACKLIST,
		Deps:        c.IsolatedHashes,
		IsolateFile: path.Join(infraBotsDir, c.TaskSpec.Isolate),
		OsType:      os,
	}
}

// getPatchStorage returns "gerrit" or "rietveld" based on the Server URL.
func getPatchStorage(server string) string {
	if server == "" {
		return ""
	}
	if strings.Contains(server, "codereview.chromium") {
		return "rietveld"
	}
	return "gerrit"
}

// replaceVars replaces variable names with their values in a given string.
func replaceVars(c *taskCandidate, s string) string {
	issueShort := ""
	if len(c.Issue) < specs.ISSUE_SHORT_LENGTH {
		issueShort = c.Issue
	} else {
		issueShort = c.Issue[len(c.Issue)-specs.ISSUE_SHORT_LENGTH:]
	}
	replacements := map[string]string{
		specs.VARIABLE_CODEREVIEW_SERVER: c.Server,
		specs.VARIABLE_ISSUE:             c.Issue,
		specs.VARIABLE_ISSUE_SHORT:       issueShort,
		specs.VARIABLE_PATCH_STORAGE:     getPatchStorage(c.Server),
		specs.VARIABLE_PATCHSET:          c.Patchset,
		specs.VARIABLE_REPO:              c.Repo,
		specs.VARIABLE_REVISION:          c.Revision,
		specs.VARIABLE_TASK_NAME:         c.Name,
	}
	for k, v := range replacements {
		s = strings.Replace(s, fmt.Sprintf(specs.VARIABLE_SYNTAX, k), v, -1)
	}
	return s
}

// MakeTaskRequest creates a SwarmingRpcsNewTaskRequest object from the taskCandidate.
func (c *taskCandidate) MakeTaskRequest(id string) *swarming_api.SwarmingRpcsNewTaskRequest {
	var cipdInput *swarming_api.SwarmingRpcsCipdInput
	if len(c.TaskSpec.CipdPackages) > 0 {
		cipdInput = &swarming_api.SwarmingRpcsCipdInput{
			Packages: make([]*swarming_api.SwarmingRpcsCipdPackage, 0, len(c.TaskSpec.CipdPackages)),
		}
		for _, p := range c.TaskSpec.CipdPackages {
			cipdInput.Packages = append(cipdInput.Packages, &swarming_api.SwarmingRpcsCipdPackage{
				PackageName: p.Name,
				Path:        p.Path,
				Version:     p.Version,
			})
		}
	}

	dims := make([]*swarming_api.SwarmingRpcsStringPair, 0, len(c.TaskSpec.Dimensions))
	dimsMap := make(map[string]string, len(c.TaskSpec.Dimensions))
	for _, d := range c.TaskSpec.Dimensions {
		split := strings.SplitN(d, ":", 2)
		key := split[0]
		val := split[1]
		dims = append(dims, &swarming_api.SwarmingRpcsStringPair{
			Key:   key,
			Value: val,
		})
		dimsMap[key] = val
	}

	var env []*swarming_api.SwarmingRpcsStringPair
	if len(c.TaskSpec.Environment) > 0 {
		env = make([]*swarming_api.SwarmingRpcsStringPair, 0, len(c.TaskSpec.Environment))
		for k, v := range c.TaskSpec.Environment {
			env = append(env, &swarming_api.SwarmingRpcsStringPair{
				Key:   k,
				Value: v,
			})
		}
	}

	extraArgs := make([]string, 0, len(c.TaskSpec.ExtraArgs))
	for _, arg := range c.TaskSpec.ExtraArgs {
		extraArgs = append(extraArgs, replaceVars(c, arg))
	}

	expirationSecs := int64(c.TaskSpec.Expiration.Seconds())
	if expirationSecs == int64(0) {
		expirationSecs = int64(swarming.RECOMMENDED_EXPIRATION.Seconds())
	}
	executionTimeoutSecs := int64(c.TaskSpec.ExecutionTimeout.Seconds())
	if executionTimeoutSecs == int64(0) {
		executionTimeoutSecs = int64(swarming.RECOMMENDED_HARD_TIMEOUT.Seconds())
	}
	ioTimeoutSecs := int64(c.TaskSpec.IoTimeout.Seconds())
	if ioTimeoutSecs == int64(0) {
		ioTimeoutSecs = int64(swarming.RECOMMENDED_IO_TIMEOUT.Seconds())
	}
	return &swarming_api.SwarmingRpcsNewTaskRequest{
		ExpirationSecs: expirationSecs,
		Name:           c.Name,
		Priority:       int64(100.0 * c.TaskSpec.Priority),
		Properties: &swarming_api.SwarmingRpcsTaskProperties{
			CipdInput:            cipdInput,
			Dimensions:           dims,
			Env:                  env,
			ExecutionTimeoutSecs: executionTimeoutSecs,
			ExtraArgs:            extraArgs,
			InputsRef: &swarming_api.SwarmingRpcsFilesRef{
				Isolated:       c.IsolatedInput,
				Isolatedserver: isolate.ISOLATE_SERVER_URL,
				Namespace:      isolate.DEFAULT_NAMESPACE,
			},
			IoTimeoutSecs: ioTimeoutSecs,
		},
		Tags: db.TagsForTask(c.Name, id, c.TaskSpec.Priority, c.RepoState, c.RetryOf, dimsMap, c.ForcedJobId, c.ParentTaskIds),
		User: "skia-task-scheduler",
	}
}

// allDepsMet determines whether all dependencies for the given task candidate
// have been satisfied, and if so, returns a map of whose keys are task IDs and
// values are their isolated outputs.
func (c *taskCandidate) allDepsMet(cache db.TaskCache) (bool, map[string]string, error) {
	rv := make(map[string]string, len(c.TaskSpec.Dependencies))
	for _, depName := range c.TaskSpec.Dependencies {
		key := c.TaskKey.Copy()
		key.Name = depName
		byKey, err := cache.GetTasksByKey(&key)
		if err != nil {
			return false, nil, err
		}
		ok := false
		for _, t := range byKey {
			if t.Done() && t.Success() && t.IsolatedOutput != "" {
				rv[t.Id] = t.IsolatedOutput
				ok = true
				break
			}
		}
		if !ok {
			return false, nil, nil
		}
	}
	return true, rv, nil
}

// taskCandidateSlice is an alias used for sorting a slice of taskCandidates.
type taskCandidateSlice []*taskCandidate

func (s taskCandidateSlice) Len() int { return len(s) }
func (s taskCandidateSlice) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}
func (s taskCandidateSlice) Less(i, j int) bool {
	return s[i].Score > s[j].Score // candidates sort in decreasing order.
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"time"

	"github.com/grokify/gotilla/fmt/fmtutil"
	"github.com/grokify/gotilla/time/timeutil"
)

func getIncident(t0s, t1s, format string, durRange time.Duration, impactNum, totalNum int) {
	durPct := timeutil.DurationPct{DurationRange: durRange}
	t0, err := time.Parse(format, t0s)
	if err != nil {
		panic(err)
	}
	t1, err := time.Parse(format, t1s)
	if err != nil {
		panic(err)
	}
	durPct.DurationStartTime = t0
	durPct.DurationEndTime = t1

	fmt.Println(durPct.DurationRange.String())

	impPct := timeutil.ImpactPct{
		ImpactNum: impactNum,
		TotalNum:  totalNum}

	inc := timeutil.Event{
		DurationPct: durPct,
		ImpactPct:   impPct}
	err = inc.Inflate()
	if err != nil {
		panic(err)
	}

	fmtutil.PrintJSON(inc)

}

func main() {
	dt, err := time.Parse(time.RFC3339, "2017-08-08T00:00:00Z")
	if err != nil {
		panic(err)
	}
	rangeDur := timeutil.QuarterDuration(dt)
	fmt.Println(rangeDur.String())

	getIncident(
		"2017-09-04T09:15:00-0700",
		"2017-09-05T12:11:04-0600",
		timeutil.ISO8601Z, rangeDur, 140, 6577)

	getIncident(
		"2017-09-29T17:09:00Z",
		"2017-10-02T07:02:00Z",
		time.RFC3339, rangeDur, 180, 6577)

	fmt.Println("DONE")
}

/*
func getInc1(t0, t1 string, format string, rangeDur time.Duration, impactNum int) (timeutil.Incident, error) {
	t10, err := time.Parse(format, t0)
	if err != nil {
		return timeutil.Incident{}, err
	}
	t11, err := time.Parse(format, t1)
	if err != nil {
		return timeutil.Incident{}, err
	}
	inc1 := timeutil.Incident{
		StartTime:     t10,
		EndTime:       t11,
		ImpactNum:     impactNum,
		TotalNum:      6577,
		RangeDuration: rangeDur,
	}

	err = inc1.Inflate()
	if err != nil {
		return timeutil.Incident{}, err
	}

	return inc1, nil
}

func main() {
	dt, err := time.Parse(time.RFC3339, "2017-08-08T00:00:00Z")
	if err != nil {
		panic(err)
	}
	rangeDur, err := timeutil.QuarterDuration(dt)
	if err != nil {
		panic(err)
	}
	fmt.Println(rangeDur.String())

	inc1, err := getInc1(
		"2017-09-04T09:15:00-0700",
		"2017-09-05T12:11:04-0600",
		timeutil.ISO8601Z4, rangeDur, 142)
	if err != nil {
		panic(err)
	}
	fmtutil.PrintJSON(inc1)
	fmt.Println(inc1.Duration.String())

	inc2, err := getInc1(
		"2017-09-29T17:09:00Z",
		"2017-10-02T07:02:00Z",
		time.RFC3339, rangeDur, 180)
	if err != nil {
		panic(err)
	}
	fmtutil.PrintJSON(inc2)
	fmt.Println(inc2.Duration.String())

	fmt.Println("DONE")
}

*/
<commit_msg>update uptime example<commit_after>package main

import (
	"fmt"
	"time"

	"github.com/grokify/gotilla/fmt/fmtutil"
	"github.com/grokify/gotilla/time/timeutil"
)

func getIncident(t0s, t1s, format string, durRange time.Duration, impactNum, totalNum int) {
	durPct := timeutil.DurationPct{DurationRange: durRange}
	t0, err := time.Parse(format, t0s)
	if err != nil {
		panic(err)
	}
	t1, err := time.Parse(format, t1s)
	if err != nil {
		panic(err)
	}
	durPct.DurationStartTime = t0
	durPct.DurationEndTime = t1

	fmt.Println(durPct.DurationRange.String())

	impPct := timeutil.ImpactPct{
		ImpactNum: impactNum,
		TotalNum:  totalNum}

	inc := timeutil.Event{
		DurationPct: durPct,
		ImpactPct:   impPct}
	err = inc.Inflate()
	if err != nil {
		panic(err)
	}

	fmtutil.PrintJSON(inc)

}

func main() {
	dt, err := time.Parse(time.RFC3339, "2017-08-08T00:00:00Z")
	if err != nil {
		panic(err)
	}
	rangeDur := timeutil.QuarterDuration(dt)
	fmt.Println(rangeDur.String())

	getIncident(
		"2017-09-04T09:15:00-0700",
		"2017-09-05T12:11:04-0600",
		timeutil.ISO8601, rangeDur, 140, 6577)

	getIncident(
		"2017-09-29T17:09:00Z",
		"2017-10-02T07:02:00Z",
		time.RFC3339, rangeDur, 180, 6577)

	fmt.Println("DONE")
}

/*
func getInc1(t0, t1 string, format string, rangeDur time.Duration, impactNum int) (timeutil.Incident, error) {
	t10, err := time.Parse(format, t0)
	if err != nil {
		return timeutil.Incident{}, err
	}
	t11, err := time.Parse(format, t1)
	if err != nil {
		return timeutil.Incident{}, err
	}
	inc1 := timeutil.Incident{
		StartTime:     t10,
		EndTime:       t11,
		ImpactNum:     impactNum,
		TotalNum:      6577,
		RangeDuration: rangeDur,
	}

	err = inc1.Inflate()
	if err != nil {
		return timeutil.Incident{}, err
	}

	return inc1, nil
}

func main() {
	dt, err := time.Parse(time.RFC3339, "2017-08-08T00:00:00Z")
	if err != nil {
		panic(err)
	}
	rangeDur, err := timeutil.QuarterDuration(dt)
	if err != nil {
		panic(err)
	}
	fmt.Println(rangeDur.String())

	inc1, err := getInc1(
		"2017-09-04T09:15:00-0700",
		"2017-09-05T12:11:04-0600",
		timeutil.ISO8601Z4, rangeDur, 142)
	if err != nil {
		panic(err)
	}
	fmtutil.PrintJSON(inc1)
	fmt.Println(inc1.Duration.String())

	inc2, err := getInc1(
		"2017-09-29T17:09:00Z",
		"2017-10-02T07:02:00Z",
		time.RFC3339, rangeDur, 180)
	if err != nil {
		panic(err)
	}
	fmtutil.PrintJSON(inc2)
	fmt.Println(inc2.Duration.String())

	fmt.Println("DONE")
}

*/
<|endoftext|>