<commit_before>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/daviddengcn/go-colortext"
	"github.com/korovkin/limiter"
	"github.com/korovkin/parallel"
)

type logger struct {
	ticket   int
	hostname string
	buf      *bytes.Buffer
}

var (
	loggerMutex     = new(sync.Mutex)
	loggerIndex     = int(0)
	loggerStartTime = time.Now()
	loggerHostname  = ""
)

var loggerColors = []ct.Color{
	ct.Green,
	ct.Cyan,
	ct.Magenta,
	ct.Yellow,
	ct.Blue,
	ct.Red,
}

func (l *logger) Write(p []byte) (int, error) {
	buf := bytes.NewBuffer(p)
	wrote := 0
	for {
		line, err := buf.ReadBytes('\n')
		if len(line) > 1 {
			now := time.Now().Format("15:04:05")
			s := string(line)
			ts := time.Since(loggerStartTime).String()

			loggerMutex.Lock()
			ct.ChangeColor(loggerColors[l.ticket%len(loggerColors)], false, ct.None, false)
			fmt.Printf("[%16s %s %s %d] ", ts, l.hostname, now, l.ticket)
			ct.ResetColor()

			if l.buf != nil {
				l.buf.Write([]byte(s))
			}

			fmt.Print(s)
			loggerMutex.Unlock()

			wrote += len(line)
		}
		if err != nil {
			break
		}
	}
	if len(p) > 0 && p[len(p)-1] != '\n' {
		fmt.Println()
	}

	return len(p), nil
}

func newLogger(ticket int, collectLines bool) *logger {
	l := &logger{ticket: ticket, buf: nil}
	if collectLines {
		l.buf = &bytes.Buffer{}
	}
	return l
}

func executeCommand(p *Parallel, ticket int, cmdLine string) (*parallel.Output, error) {
	T_START := time.Now()
	var err error
	output := &parallel.Output{}
	loggerOut := newLogger(ticket, true)
	loggerErr := newLogger(ticket, true)

	defer func() {
		fmt.Fprintf(
			loggerOut,
			"execute: done: dt: "+time.Since(T_START).String()+"\n",
		)
	}()

	// execute remotely:
	if len(p.Slaves) > 0 {
		slave := p.Slaves[ticket%len(p.Slaves)]

		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to GetTransport:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()
		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)

		cmd := parallel.Cmd{
			CmdLine: cmdLine,
			Ticket:  int64(ticket),
		}

		output, err = client.Execute(&cmd)
		if err != nil {
			log.Fatalln("failed to execute:", err.Error())
		}

		hostname := output.Tags["hostname"]
		loggerOut.hostname = hostname
		loggerErr.hostname = hostname

		fmt.Fprintf(loggerOut,
			"execute: remotely: host: %s stdout: [%s]\n",
			hostname,
			output.Stdout)

		fmt.Fprintf(loggerErr,
			"execute: remotely: host: %s stderr: [%s]\n",
			hostname,
			output.Stderr)

		return output, err
	}

	// execute locally:
	cs := []string{"/bin/sh", "-c", cmdLine}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = nil
	cmd.Stdout = loggerOut
	cmd.Stderr = loggerErr
	cmd.Env = append(
		os.Environ(),
		fmt.Sprintf("PARALLEL_TICKER=%d", ticket),
	)

	fmt.Fprintf(loggerOut, "run: '"+cmdLine+"'\n")

	err = cmd.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
		return output, err
	}

	if err == nil {
		err = cmd.Wait()
	}

	output.Tags = map[string]string{"hostname": loggerHostname}

	if loggerOut.buf != nil {
		output.Stdout = string(loggerOut.buf.Bytes())
	}

	if loggerErr.buf != nil {
		output.Stderr = string(loggerErr.buf.Bytes())
	}

	return output, err
}

type Slave struct {
	Address string `json:"address"`
}

type Parallel struct {
	jobs   int
	logger *logger
	worker *limiter.ConcurrencyLimiter

	// master / slave
	protocolFactory  thrift.TProtocolFactory
	transportFactory thrift.TTransportFactory

	// master:
	Slaves []*Slave

	// slave:
	handler         *ParallelSlaveHandler
	serverTransport thrift.TServerTransport
	address         string
}

func (p *Parallel) Close() {
	p.worker.Wait()
}

func mainMaster(p *Parallel) {
	var err error

	// connect to slaves:
	for _, slave := range p.Slaves {
		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()

		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)
		ok, err := client.Ping()
		if err != nil {
			log.Fatalln("failed to ping client:", err.Error())
		}

		fmt.Fprintf(p.logger, fmt.Sprintf("adding slave: %s ok: %s", slave.Address, ok))

	}

	r := bufio.NewReaderSize(os.Stdin, 1*1024*1024)
	fmt.Fprintf(p.logger, "reading from stdin...\n")
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)

		p.worker.ExecuteWithTicket(func(ticket int) {
			executeCommand(p, ticket, line)
		})
	}
}

type ParallelSlaveHandler struct {
	p *Parallel
}

func NewParallelSlaveHandler() *ParallelSlaveHandler {
	return &ParallelSlaveHandler{}
}

func (p *ParallelSlaveHandler) Execute(command *parallel.Cmd) (output *parallel.Output, err error) {
	output = nil
	err = nil
	output, err = executeCommand(p.p, int(command.Ticket), command.CmdLine)

	// TODO:: recover, handle panics

	return output, err
}

func (p *ParallelSlaveHandler) Ping() (r string, err error) {
	return "ping:ok", nil
}

func mainSlave(p *Parallel) {
	var err error

	p.serverTransport, err = thrift.NewTServerSocket(p.address)

	if err != nil {
		log.Fatalln("failed to start server:", err.Error())
		return
	}

	p.handler = NewParallelSlaveHandler()
	p.handler.p = p

	server := thrift.NewTSimpleServer4(
		parallel.NewParallelProcessor(p.handler),
		p.serverTransport,
		p.transportFactory,
		p.protocolFactory)

	err = server.Serve()
	if err != nil {
		log.Fatalln("failed to run slave:", err.Error())
	}
}

func main() {
	T_START := time.Now()
	logger := newLogger(0, false)
	defer func() {
		fmt.Fprintf(logger, "all done: dt: "+time.Since(T_START).String()+"\n")
	}()

	flag_jobs := flag.Int(
		"j",
		2,
		"num of concurrent jobs")

	flag_slave := flag.Bool(
		"slave",
		false,
		"run as slave")

	flag_slaves := flag.String(
		"slaves",
		"",
		"CSV list of slave addresses")

	address := flag.String(
		"address",
		"localhost:9010",
		"slave address")

	loggerHostname, _ = os.Hostname()

	flag.Parse()
	fmt.Fprintf(logger, fmt.Sprintf("concurrency limit: %d", *flag_jobs))
	fmt.Fprintf(logger, fmt.Sprintf("slaves: %s", *flag_slaves))

	p := Parallel{}
	p.jobs = *flag_jobs
	p.logger = logger
	p.worker = limiter.NewConcurrencyLimiter(p.jobs)
	p.address = *address
	p.Slaves = []*Slave{}
	for _, slaveAddr := range strings.Split(*flag_slaves, ",") {
		if slaveAddr != "" {
			slave := Slave{Address: slaveAddr}
			p.Slaves = append(p.Slaves, &slave)
		}
	}

	defer p.Close()

	// thrift protocol
	p.protocolFactory = thrift.NewTBinaryProtocolFactoryDefault()

	// thrift transport
	p.transportFactory = thrift.NewTTransportFactory()
	p.transportFactory = thrift.NewTFramedTransportFactory(p.transportFactory)

	if *flag_slave == false {
		loggerHostname = p.address
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, fmt.Sprintf("running as master\n"))
		mainMaster(&p)
	} else {
		loggerHostname = p.address
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, fmt.Sprintf("running as slave on: %s\n", p.address))
		mainSlave(&p)
	}
}
<commit_msg>cleanup<commit_after>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/daviddengcn/go-colortext"
	"github.com/korovkin/limiter"
	"github.com/korovkin/parallel"
)

type logger struct {
	ticket   int
	hostname string
	buf      *bytes.Buffer
}

var (
	loggerMutex     = new(sync.Mutex)
	loggerIndex     = int(0)
	loggerStartTime = time.Now()
	loggerHostname  = ""
)

var loggerColors = []ct.Color{
	ct.Green,
	ct.Cyan,
	ct.Magenta,
	ct.Yellow,
	ct.Blue,
	ct.Red,
}

func (l *logger) Write(p []byte) (int, error) {
	buf := bytes.NewBuffer(p)
	wrote := 0
	for {
		line, err := buf.ReadBytes('\n')
		if len(line) > 1 {
			now := time.Now().Format("15:04:05")
			s := string(line)
			ts := time.Since(loggerStartTime).String()

			loggerMutex.Lock()
			ct.ChangeColor(loggerColors[l.ticket%len(loggerColors)], false, ct.None, false)
			fmt.Printf("[%16s %s %s %d] ", ts, l.hostname, now, l.ticket)
			ct.ResetColor()

			if l.buf != nil {
				l.buf.Write([]byte(s))
			}

			fmt.Print(s)
			loggerMutex.Unlock()

			wrote += len(line)
		}
		if err != nil {
			break
		}
	}
	if len(p) > 0 && p[len(p)-1] != '\n' {
		fmt.Println()
	}

	return len(p), nil
}

func newLogger(ticket int, collectLines bool) *logger {
	l := &logger{ticket: ticket, buf: nil}
	if collectLines {
		l.buf = &bytes.Buffer{}
	}
	return l
}

func executeCommand(p *Parallel, ticket int, cmdLine string) (*parallel.Output, error) {
	T_START := time.Now()
	var err error
	output := &parallel.Output{}
	loggerOut := newLogger(ticket, true)
	loggerErr := newLogger(ticket, true)

	defer func() {
		fmt.Fprintf(
			loggerOut,
			"execute: done: dt: "+time.Since(T_START).String()+"\n",
		)
	}()

	// execute remotely:
	if len(p.Slaves) > 0 {
		slave := p.Slaves[ticket%len(p.Slaves)]

		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to GetTransport:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()
		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)

		cmd := parallel.Cmd{
			CmdLine: cmdLine,
			Ticket:  int64(ticket),
		}

		output, err = client.Execute(&cmd)
		if err != nil {
			log.Fatalln("failed to execute:", err.Error())
		}

		hostname := output.Tags["hostname"]
		loggerOut.hostname = hostname
		loggerErr.hostname = hostname

		fmt.Fprintf(loggerOut,
			"execute: remotely: host: %s stdout: [%s]\n",
			hostname,
			output.Stdout)

		fmt.Fprintf(loggerErr,
			"execute: remotely: host: %s stderr: [%s]\n",
			hostname,
			output.Stderr)

		return output, err
	}

	// execute locally:
	cs := []string{"/bin/sh", "-c", cmdLine}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = nil
	cmd.Stdout = loggerOut
	cmd.Stderr = loggerErr
	cmd.Env = append(
		os.Environ(),
		fmt.Sprintf("PARALLEL_TICKER=%d", ticket),
	)

	fmt.Fprintf(loggerOut, "run: '"+cmdLine+"'\n")

	err = cmd.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
		return output, err
	}

	if err == nil {
		err = cmd.Wait()
	}

	output.Tags = map[string]string{"hostname": loggerHostname}

	if loggerOut.buf != nil {
		output.Stdout = string(loggerOut.buf.Bytes())
	}

	if loggerErr.buf != nil {
		output.Stderr = string(loggerErr.buf.Bytes())
	}

	return output, err
}

type Slave struct {
	Address string `json:"address"`
}

type Parallel struct {
	jobs   int
	logger *logger
	worker *limiter.ConcurrencyLimiter

	// master / slave
	protocolFactory  thrift.TProtocolFactory
	transportFactory thrift.TTransportFactory

	// master:
	Slaves []*Slave

	// slave:
	handler         *ParallelSlaveHandler
	serverTransport thrift.TServerTransport
	slaveAddress    string
}

func (p *Parallel) Close() {
	p.worker.Wait()
}

func mainMaster(p *Parallel) {
	var err error

	// connect to slaves:
	for _, slave := range p.Slaves {
		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport, err = p.transportFactory.GetTransport(transport)
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()

		client := parallel.NewParallelClientFactory(transport, p.protocolFactory)
		ok, err := client.Ping()
		if err != nil {
			log.Fatalln("failed to ping client:", err.Error())
		}

		fmt.Fprintf(p.logger, fmt.Sprintf("adding slave: %s ok: %s", slave.Address, ok))

	}

	r := bufio.NewReaderSize(os.Stdin, 1*1024*1024)
	fmt.Fprintf(p.logger, "reading from stdin...\n")
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)

		p.worker.ExecuteWithTicket(func(ticket int) {
			executeCommand(p, ticket, line)
		})
	}
}

type ParallelSlaveHandler struct {
	p *Parallel
}

func NewParallelSlaveHandler() *ParallelSlaveHandler {
	return &ParallelSlaveHandler{}
}

func (p *ParallelSlaveHandler) Execute(command *parallel.Cmd) (output *parallel.Output, err error) {
	output = nil
	err = nil
	output, err = executeCommand(p.p, int(command.Ticket), command.CmdLine)

	// TODO:: recover, handle panics

	return output, err
}

func (p *ParallelSlaveHandler) Ping() (r string, err error) {
	return "ping:ok", nil
}

func mainSlave(p *Parallel) {
	var err error

	p.serverTransport, err = thrift.NewTServerSocket(p.slaveAddress)

	if err != nil {
		log.Fatalln("failed to start server:", err.Error())
		return
	}

	p.handler = NewParallelSlaveHandler()
	p.handler.p = p

	server := thrift.NewTSimpleServer4(
		parallel.NewParallelProcessor(p.handler),
		p.serverTransport,
		p.transportFactory,
		p.protocolFactory)

	err = server.Serve()
	if err != nil {
		log.Fatalln("failed to run slave:", err.Error())
	}
}

func main() {
	T_START := time.Now()
	logger := newLogger(0, false)
	defer func() {
		fmt.Fprintf(logger, "all done: dt: "+time.Since(T_START).String()+"\n")
	}()

	flag_jobs := flag.Int(
		"j",
		2,
		"num of concurrent jobs")

	flag_slave := flag.Bool(
		"slave",
		false,
		"run as slave")

	flag_slaves := flag.String(
		"slaves",
		"",
		"CSV list of slave addresses")

	flag_slave_address := flag.String(
		"address",
		"localhost:9010",
		"slave address")

	loggerHostname, _ = os.Hostname()

	flag.Parse()
	fmt.Fprintf(logger, fmt.Sprintf("concurrency limit: %d", *flag_jobs))
	fmt.Fprintf(logger, fmt.Sprintf("slaves: %s", *flag_slaves))

	p := Parallel{}
	p.jobs = *flag_jobs
	p.logger = logger
	p.worker = limiter.NewConcurrencyLimiter(p.jobs)
	p.slaveAddress = *flag_slave_address
	p.Slaves = []*Slave{}
	for _, slaveAddr := range strings.Split(*flag_slaves, ",") {
		if slaveAddr != "" {
			slave := Slave{Address: slaveAddr}
			p.Slaves = append(p.Slaves, &slave)
		}
	}

	defer p.Close()

	// thrift protocol
	p.protocolFactory = thrift.NewTBinaryProtocolFactoryDefault()

	// thrift transport
	p.transportFactory = thrift.NewTTransportFactory()
	p.transportFactory = thrift.NewTFramedTransportFactory(p.transportFactory)

	if *flag_slave == false {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as master\n")
		mainMaster(&p)
	} else {
		loggerHostname = p.slaveAddress
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, "running as slave on: %s\n", p.slaveAddress)
		mainSlave(&p)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/daviddengcn/go-colortext"
	"github.com/korovkin/limiter"
	"github.com/korovkin/parallel"
)

type logger struct {
	ticket int
	buf    *bytes.Buffer
}

var (
	loggerMutex     = new(sync.Mutex)
	loggerIndex     = int(0)
	loggerStartTime = time.Now()
	loggerHostname  = ""
)

var loggerColors = []ct.Color{
	ct.Green,
	ct.Cyan,
	ct.Magenta,
	ct.Yellow,
	ct.Blue,
	ct.Red,
}

func (l *logger) Write(p []byte) (int, error) {
	buf := bytes.NewBuffer(p)
	wrote := 0
	for {
		line, err := buf.ReadBytes('\n')
		if len(line) > 1 {
			now := time.Now().Format("15:04:05")
			s := string(line)
			ts := time.Since(loggerStartTime).String()

			loggerMutex.Lock()
			ct.ChangeColor(loggerColors[l.ticket%len(loggerColors)], false, ct.None, false)
			fmt.Printf("[%16s %s %s %d] ", ts, loggerHostname, now, l.ticket)
			ct.ResetColor()

			if l.buf != nil {
				l.buf.Write([]byte(s))
			}

			fmt.Print(s)
			loggerMutex.Unlock()

			wrote += len(line)
		}
		if err != nil {
			break
		}
	}
	if len(p) > 0 && p[len(p)-1] != '\n' {
		fmt.Println()
	}

	return len(p), nil
}

func newLogger(ticket int, collectLines bool) *logger {
	l := &logger{ticket: ticket, buf: nil}
	if collectLines {
		l.buf = &bytes.Buffer{}
	}
	return l
}

func executeCommand(p *Parallel, ticket int, cmdLine string) (*parallel.Output, error) {
	T_START := time.Now()
	var err error
	output := &parallel.Output{}

	loggerOut := newLogger(ticket, true)
	loggerErr := newLogger(ticket, true)

	defer func() {
		fmt.Fprintf(
			loggerOut,
			"execute: done: dt: "+time.Since(T_START).String()+"\n",
		)
	}()

	// execute remotely:
	if len(p.Slaves) > 0 {
		slave := p.Slaves[ticket%len(p.Slaves)]

		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport = p.transportFactory.GetTransport(transport)
		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()
		slave.Client = parallel.NewParallelClientFactory(transport, p.protocolFactory)

		output, err = slave.Client.Execute(&parallel.Cmd{
			CmdLine: cmdLine,
			Ticket:  int64(ticket),
		})
		if err != nil {
			log.Fatalln("failed to execute:", err.Error())
		}

		hostname := output.Tags["hostname"]

		fmt.Fprintf(loggerOut, "execute: remotely: host: %s stdout: '%s'\n", hostname, output.Stdout)
		fmt.Fprintf(loggerErr, "execute: remotely: host: %s stderr: '%s'\n", hostname, output.Stderr)

		return output, err
	}

	// execute locally:
	cs := []string{"/bin/sh", "-c", cmdLine}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = nil
	cmd.Stdout = loggerOut
	cmd.Stderr = loggerErr
	cmd.Env = append(
		os.Environ(),
		fmt.Sprintf("PARALLEL_TICKER=%d", ticket),
	)

	fmt.Fprintf(loggerOut, "run: '"+cmdLine+"'\n")

	err = cmd.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
		return output, err
	}

	if err == nil {
		err = cmd.Wait()
	}

	output.Tags = map[string]string{"hostname": loggerHostname}

	if loggerOut.buf != nil {
		output.Stdout = string(loggerOut.buf.Bytes())
	}

	if loggerErr.buf != nil {
		output.Stderr = string(loggerErr.buf.Bytes())
	}

	return output, err
}

type Slave struct {
	Address string `json:"address"`
	Client  *parallel.ParallelClient
}

type Parallel struct {
	jobs   int
	logger *logger
	worker *limiter.ConcurrencyLimiter

	// master / slave
	protocolFactory  thrift.TProtocolFactory
	transportFactory thrift.TTransportFactory

	// master:
	Slaves []*Slave

	// slave:
	handler         *ParallelSlaveHandler
	serverTransport thrift.TServerTransport
	address         string
}

func (p *Parallel) Close() {
	p.worker.Wait()
}

func mainMaster(p *Parallel) {
	var err error

	// connect to slaves:
	for _, slave := range p.Slaves {
		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport = p.transportFactory.GetTransport(transport)
		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}
		defer transport.Close()

		slave.Client = parallel.NewParallelClientFactory(transport, p.protocolFactory)
		ok, err := slave.Client.Ping()
		if err != nil {
			log.Fatalln("failed to ping client:", err.Error())
		}

		fmt.Fprintf(p.logger, fmt.Sprintf("adding slave: %s ok: %s", slave.Address, ok))

	}

	r := bufio.NewReaderSize(os.Stdin, 1*1024*1024)
	fmt.Fprintf(p.logger, "reading from stdin...\n")
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)

		p.worker.ExecuteWithTicket(func(ticket int) {
			executeCommand(p, ticket, line)
		})
	}
}

type ParallelSlaveHandler struct {
	p *Parallel
}

func NewParallelSlaveHandler() *ParallelSlaveHandler {
	return &ParallelSlaveHandler{}
}

func (p *ParallelSlaveHandler) Execute(command *parallel.Cmd) (output *parallel.Output, err error) {
	log.Println("ParallelSlaveHandler: Execute: ", command.CmdLine)
	output, err = executeCommand(p.p, int(command.Ticket), command.CmdLine)
	// TODO:: recover, handle panics
	return output, err
}

func (p *ParallelSlaveHandler) Ping() (r string, err error) {
	log.Println("ParallelSlaveHandler: Ping")
	return "ping:ok", nil
}

func mainSlave(p *Parallel) {
	var err error

	p.serverTransport, err = thrift.NewTServerSocket(p.address)

	if err != nil {
		log.Fatalln("failed to start server:", err.Error())
		return
	}

	p.handler = NewParallelSlaveHandler()
	p.handler.p = p

	server := thrift.NewTSimpleServer4(
		parallel.NewParallelProcessor(p.handler),
		p.serverTransport,
		p.transportFactory,
		p.protocolFactory)

	err = server.Serve()
	if err != nil {
		log.Fatalln("failed to run slave:", err.Error())
	}
}

func main() {
	T_START := time.Now()
	logger := newLogger(0, false)
	defer func() {
		fmt.Fprintf(logger, "all done: dt: "+time.Since(T_START).String()+"\n")
	}()

	flag_jobs := flag.Int(
		"j",
		2,
		"num of concurrent jobs")

	flag_slave := flag.Bool(
		"slave",
		false,
		"run as slave")

	slaves := flag.String(
		"slaves",
		"",
		"CSV list of slave addresses")

	address := flag.String(
		"address",
		"localhost:9010",
		"slave address")

	loggerHostname, _ = os.Hostname()

	flag.Parse()
	fmt.Fprintf(logger, fmt.Sprintf("concurrency limit: %d", *flag_jobs))
	fmt.Fprintf(logger, fmt.Sprintf("slaves: %s", *slaves))

	p := Parallel{}
	p.jobs = *flag_jobs
	p.logger = logger
	p.worker = limiter.NewConcurrencyLimiter(p.jobs)
	p.address = *address
	p.Slaves = []*Slave{}
	for _, slaveAddr := range strings.Split(*slaves, ",") {
		if slaveAddr == "" {
			continue
		}
		slave := Slave{Address: slaveAddr}
		p.Slaves = append(p.Slaves, &slave)
	}

	defer p.Close()

	// thrift protocol
	p.protocolFactory = thrift.NewTBinaryProtocolFactoryDefault()

	// thrift transport
	p.transportFactory = thrift.NewTTransportFactory()
	p.transportFactory = thrift.NewTFramedTransportFactory(p.transportFactory)

	if *flag_slave == false {
		fmt.Fprintf(logger, fmt.Sprintf("running as master\n"))
		mainMaster(&p)
	} else {
		fmt.Fprintf(logger, fmt.Sprintf("running as slave on: %s\n", p.address))
		mainSlave(&p)
	}
}
<commit_msg>print remote hostname<commit_after>package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"git.apache.org/thrift.git/lib/go/thrift"
	"github.com/daviddengcn/go-colortext"
	"github.com/korovkin/limiter"
	"github.com/korovkin/parallel"
)

type logger struct {
	ticket   int
	hostname string
	buf      *bytes.Buffer
}

var (
	loggerMutex     = new(sync.Mutex)
	loggerIndex     = int(0)
	loggerStartTime = time.Now()
	loggerHostname  = ""
)

var loggerColors = []ct.Color{
	ct.Green,
	ct.Cyan,
	ct.Magenta,
	ct.Yellow,
	ct.Blue,
	ct.Red,
}

func (l *logger) Write(p []byte) (int, error) {
	buf := bytes.NewBuffer(p)
	wrote := 0
	for {
		line, err := buf.ReadBytes('\n')
		if len(line) > 1 {
			now := time.Now().Format("15:04:05")
			s := string(line)
			ts := time.Since(loggerStartTime).String()

			loggerMutex.Lock()
			ct.ChangeColor(loggerColors[l.ticket%len(loggerColors)], false, ct.None, false)
			fmt.Printf("[%16s %s %s %d] ", ts, l.hostname, now, l.ticket)
			ct.ResetColor()

			if l.buf != nil {
				l.buf.Write([]byte(s))
			}

			fmt.Print(s)
			loggerMutex.Unlock()

			wrote += len(line)
		}
		if err != nil {
			break
		}
	}
	if len(p) > 0 && p[len(p)-1] != '\n' {
		fmt.Println()
	}

	return len(p), nil
}

func newLogger(ticket int, collectLines bool) *logger {
	l := &logger{ticket: ticket, buf: nil}
	if collectLines {
		l.buf = &bytes.Buffer{}
	}
	return l
}

func executeCommand(p *Parallel, ticket int, cmdLine string) (*parallel.Output, error) {
	T_START := time.Now()
	var err error
	output := &parallel.Output{}

	loggerOut := newLogger(ticket, true)
	loggerErr := newLogger(ticket, true)

	defer func() {
		fmt.Fprintf(
			loggerOut,
			"execute: done: dt: "+time.Since(T_START).String()+"\n",
		)
	}()

	// execute remotely:
	if len(p.Slaves) > 0 {
		slave := p.Slaves[ticket%len(p.Slaves)]

		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport = p.transportFactory.GetTransport(transport)
		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}

		defer transport.Close()
		slave.Client = parallel.NewParallelClientFactory(transport, p.protocolFactory)

		output, err = slave.Client.Execute(&parallel.Cmd{
			CmdLine: cmdLine,
			Ticket:  int64(ticket),
		})
		if err != nil {
			log.Fatalln("failed to execute:", err.Error())
		}

		hostname := output.Tags["hostname"]
		loggerOut.hostname = hostname
		loggerErr.hostname = hostname

		fmt.Fprintf(loggerOut, "execute: remotely: host: %s stdout: [%s]\n", hostname, output.Stdout)
		fmt.Fprintf(loggerErr, "execute: remotely: host: %s stderr: [%s]\n", hostname, output.Stderr)

		return output, err
	}

	// execute locally:
	cs := []string{"/bin/sh", "-c", cmdLine}
	cmd := exec.Command(cs[0], cs[1:]...)
	cmd.Stdin = nil
	cmd.Stdout = loggerOut
	cmd.Stderr = loggerErr
	cmd.Env = append(
		os.Environ(),
		fmt.Sprintf("PARALLEL_TICKER=%d", ticket),
	)

	fmt.Fprintf(loggerOut, "run: '"+cmdLine+"'\n")

	err = cmd.Start()
	if err != nil {
		log.Fatalln("failed to start:", err)
		return output, err
	}

	if err == nil {
		err = cmd.Wait()
	}

	output.Tags = map[string]string{"hostname": loggerHostname}

	if loggerOut.buf != nil {
		output.Stdout = string(loggerOut.buf.Bytes())
	}

	if loggerErr.buf != nil {
		output.Stderr = string(loggerErr.buf.Bytes())
	}

	return output, err
}

type Slave struct {
	Address string `json:"address"`
	Client  *parallel.ParallelClient
}

type Parallel struct {
	jobs   int
	logger *logger
	worker *limiter.ConcurrencyLimiter

	// master / slave
	protocolFactory  thrift.TProtocolFactory
	transportFactory thrift.TTransportFactory

	// master:
	Slaves []*Slave

	// slave:
	handler         *ParallelSlaveHandler
	serverTransport thrift.TServerTransport
	address         string
}

func (p *Parallel) Close() {
	p.worker.Wait()
}

func mainMaster(p *Parallel) {
	var err error

	// connect to slaves:
	for _, slave := range p.Slaves {
		var transport thrift.TTransport
		transport, err = thrift.NewTSocket(slave.Address)
		if err != nil {
			log.Fatalln("failed to dial slave:", err.Error())
		}

		transport = p.transportFactory.GetTransport(transport)
		err = transport.Open()
		if err != nil {
			log.Fatalln("failed to open:", err.Error())
		}
		defer transport.Close()

		slave.Client = parallel.NewParallelClientFactory(transport, p.protocolFactory)
		ok, err := slave.Client.Ping()
		if err != nil {
			log.Fatalln("failed to ping client:", err.Error())
		}

		fmt.Fprintf(p.logger, fmt.Sprintf("adding slave: %s ok: %s", slave.Address, ok))

	}

	r := bufio.NewReaderSize(os.Stdin, 1*1024*1024)
	fmt.Fprintf(p.logger, "reading from stdin...\n")
	for {
		line, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		line = strings.TrimSpace(line)

		p.worker.ExecuteWithTicket(func(ticket int) {
			executeCommand(p, ticket, line)
		})
	}
}

type ParallelSlaveHandler struct {
	p *Parallel
}

func NewParallelSlaveHandler() *ParallelSlaveHandler {
	return &ParallelSlaveHandler{}
}

func (p *ParallelSlaveHandler) Execute(command *parallel.Cmd) (output *parallel.Output, err error) {
	log.Println("ParallelSlaveHandler: Execute: ", command.CmdLine)
	output, err = executeCommand(p.p, int(command.Ticket), command.CmdLine)
	// TODO:: recover, handle panics
	return output, err
}

func (p *ParallelSlaveHandler) Ping() (r string, err error) {
	log.Println("ParallelSlaveHandler: Ping")
	return "ping:ok", nil
}

func mainSlave(p *Parallel) {
	var err error

	p.serverTransport, err = thrift.NewTServerSocket(p.address)

	if err != nil {
		log.Fatalln("failed to start server:", err.Error())
		return
	}

	p.handler = NewParallelSlaveHandler()
	p.handler.p = p

	server := thrift.NewTSimpleServer4(
		parallel.NewParallelProcessor(p.handler),
		p.serverTransport,
		p.transportFactory,
		p.protocolFactory)

	err = server.Serve()
	if err != nil {
		log.Fatalln("failed to run slave:", err.Error())
	}
}

func main() {
	T_START := time.Now()
	logger := newLogger(0, false)
	defer func() {
		fmt.Fprintf(logger, "all done: dt: "+time.Since(T_START).String()+"\n")
	}()

	flag_jobs := flag.Int(
		"j",
		2,
		"num of concurrent jobs")

	flag_slave := flag.Bool(
		"slave",
		false,
		"run as slave")

	slaves := flag.String(
		"slaves",
		"",
		"CSV list of slave addresses")

	address := flag.String(
		"address",
		"localhost:9010",
		"slave address")

	loggerHostname, _ = os.Hostname()

	flag.Parse()
	fmt.Fprintf(logger, fmt.Sprintf("concurrency limit: %d", *flag_jobs))
	fmt.Fprintf(logger, fmt.Sprintf("slaves: %s", *slaves))

	p := Parallel{}
	p.jobs = *flag_jobs
	p.logger = logger
	p.worker = limiter.NewConcurrencyLimiter(p.jobs)
	p.address = *address
	p.Slaves = []*Slave{}
	for _, slaveAddr := range strings.Split(*slaves, ",") {
		if slaveAddr == "" {
			continue
		}
		slave := Slave{Address: slaveAddr}
		p.Slaves = append(p.Slaves, &slave)
	}

	defer p.Close()

	// thrift protocol
	p.protocolFactory = thrift.NewTBinaryProtocolFactoryDefault()

	// thrift transport
	p.transportFactory = thrift.NewTTransportFactory()
	p.transportFactory = thrift.NewTFramedTransportFactory(p.transportFactory)

	if *flag_slave == false {
		loggerHostname = p.address
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, fmt.Sprintf("running as master\n"))
		mainMaster(&p)
	} else {
		loggerHostname = p.address
		logger.hostname = loggerHostname

		fmt.Fprintf(logger, fmt.Sprintf("running as slave on: %s\n", p.address))
		mainSlave(&p)
	}
}
<|endoftext|>


<commit_before>package client

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

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

func TestBootstrapSettings(t *testing.T) {
	file, err := ioutil.TempFile("", ".packaged-lantern.yaml")
	defer func() {
		err := os.Remove(file.Name())
		if err != nil {
			log.Errorf("Could not remove file? %v", err)
		}
	}()
	assert.True(t, err == nil, "Should not be an error")
	file.Close()

	log.Debugf("File at: %v", file.Name())
	settings := BootstrapSettings{StartupUrl: "test"}
	log.Debugf("Settings: %v", settings)

	data, er := yaml.Marshal(&settings)
	assert.True(t, er == nil, "Should not be an error")

	e := ioutil.WriteFile(file.Name(), data, 0644)
	assert.True(t, e == nil, "Should not be an error")

	path, ps, errr := readSettingsFromFile(file.Name())
	assert.Equal(t, "test", ps.StartupUrl, "Unexpected startup URL")
	assert.Equal(t, file.Name(), path, "Wrote to unexpected path")
	assert.True(t, errr == nil, "Should not be an error")

	// Now do another full round trip, writing and reading
	// Overwite local to avoid affecting actual Lantern instances
	local = file.Name()
	path, errrr := writeToDisk(ps)
	assert.True(t, errrr == nil, "Should not be an error")
	path, ps, err = readSettingsFromFile(path)
	assert.Equal(t, "test", ps.StartupUrl, "Could not read data")
	assert.Equal(t, local, path, "Wrote to unexpected path")
	assert.True(t, err == nil, "Should not be an error")

	url := "test"

	path, ps, err = readSettingsFromFile(path)

	log.Debugf("Wrote settings to: %v", path)
	assert.Equal(t, url, ps.StartupUrl, "Could not read data")
	assert.Equal(t, local, path, "Wrote to unexpected path")
	assert.True(t, err == nil, "Should not be an error")

	path, err = packagedSettingsPath()
	assert.True(t, err == nil, "Should not be an error")

	var dir string

	if runtime.GOOS == "darwin" {
		dir, err = filepath.Abs(filepath.Dir(os.Args[0]) + "/../Resources")
	} else if runtime.GOOS == "linux" {
		dir, err = filepath.Abs(filepath.Dir(os.Args[0]))
	}
	assert.True(t, err == nil, "Should not be an error")

	log.Debugf("Running in %v", dir)
	if runtime.GOOS == "darwin" {
		assert.Equal(t, dir+"/"+name, path, "Unexpected settings dir")
	} else if runtime.GOOS == "linux" {
		assert.Equal(t, dir+"/"+name, path, "Unexpected settings dir")
	}
}
<commit_msg>fix bootstrap_test.go<commit_after>package client

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

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

func TestBootstrapSettings(t *testing.T) {
	file, err := ioutil.TempFile("", ".packaged-lantern.yaml")
	defer func() {
		err := os.Remove(file.Name())
		if err != nil {
			log.Errorf("Could not remove file? %v", err)
		}
	}()
	assert.True(t, err == nil, "Should not be an error")
	file.Close()

	log.Debugf("File at: %v", file.Name())
	settings := BootstrapSettings{StartupUrl: "test"}
	log.Debugf("Settings: %v", settings)

	data, er := yaml.Marshal(&settings)
	assert.True(t, er == nil, "Should not be an error")

	e := ioutil.WriteFile(file.Name(), data, 0644)
	assert.True(t, e == nil, "Should not be an error")

	ps, errr := readSettingsFromFile(file.Name())
	assert.Equal(t, "test", ps.StartupUrl, "Unexpected startup URL")
	assert.True(t, errr == nil, "Should not be an error")

	// Now do another full round trip, writing and reading
	// Overwite local to avoid affecting actual Lantern instances
	local = file.Name()
	path, errrr := writeToDisk(ps)
	assert.True(t, errrr == nil, "Should not be an error")
	ps, err = readSettingsFromFile(path)
	assert.Equal(t, "test", ps.StartupUrl, "Could not read data")
	assert.Equal(t, local, path, "Wrote to unexpected path")
	assert.True(t, err == nil, "Should not be an error")

	url := "test"

	ps, err = readSettingsFromFile(path)

	log.Debugf("Wrote settings to: %v", path)
	assert.Equal(t, url, ps.StartupUrl, "Could not read data")
	assert.Equal(t, local, path, "Wrote to unexpected path")
	assert.True(t, err == nil, "Should not be an error")

	path, err = bootstrapPath(name)
	assert.True(t, err == nil, "Should not be an error")

	var dir string

	if runtime.GOOS == "darwin" {
		dir, err = filepath.Abs(filepath.Dir(os.Args[0]) + "/../Resources")
	} else if runtime.GOOS == "linux" {
		dir, err = filepath.Abs(filepath.Dir(os.Args[0]))
	}
	assert.True(t, err == nil, "Should not be an error")

	log.Debugf("Running in %v", dir)
	if runtime.GOOS == "darwin" {
		assert.Equal(t, dir+"/"+name, path, "Unexpected settings dir")
	} else if runtime.GOOS == "linux" {
		assert.Equal(t, dir+"/"+name, path, "Unexpected settings dir")
	}
}
<|endoftext|>


<commit_before>package mimic

import (
	"bytes"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"regexp"
	"testing"
	"time"

	"github.com/getlantern/testify/assert"

	proxy "github.com/getlantern/http-proxy"
)

const target = "test-data/apache-2.4.7-ubuntu14.04.raw"
const template = "test-data/apache-2.4.7-ubuntu14.04.tpl"
const current = "test-data/chained-server.raw"

type entry struct {
	method         string
	path           string
	withHostHeader bool
}

/*
For rationale behind these tests, refer following:
[1] https://github.com/getlantern/lantern/issues/1255#issuecomment-30658217
[2] https://github.com/getlantern/lantern-java/blob/master/src/main/java/org/lantern/proxy/GiveModeHttpFilters.java#L167
[3] https://github.com/apache/httpd/blob/2.4.17/server/core.c#L4379
*/
var candidates = []entry{
	// Without Host header, Apache will return 400
	{"GET", "/", false},
	{"GET", "/index.html", false},

	// 200 with default page
	{"GET", "/", true},
	{"GET", "/index.html", true},

	{"GET", "/icons/ubuntu-logo.png", true},

	// 404
	{"GET", "//cgi-bin/php", true},
	{"GET", "//cgi-bin/php5", true},
	{"GET", "//cgi-bin/php-cgi", true},
	{"GET", "//cgi-bin/php.cgi", true},
	{"GET", "//cgi-bin/php4", true},

	// multiple slashes
	{"GET", "///cgi-bin/php4", true},
	{"GET", "//cgi-bin//php4", true},

	{"GET", "/not-existed", true},
	{"GET", "/end-with-slash/", true},

	{"HEAD", "/", true},
	{"HEAD", "/index.html", true},
	{"HEAD", "/icons/ubuntu-logo.png", true},
	{"HEAD", "/not-existed", true},

	{"POST", "/", true},
	{"POST", "/index.html", true},
	{"POST", "/icons/ubuntu-logo.png", true},
	{"POST", "/not-existed", true},

	{"OPTIONS", "/", true},
	{"OPTIONS", "/index.html", true},
	{"OPTIONS", "/icons/ubuntu-logo.png", true},
	{"OPTIONS", "/not-existed", true},

	{"PUT", "/", true},
	{"PUT", "/index.html", true},
	{"PUT", "/icons/ubuntu-logo.png", true},
	{"PUT", "/not-existed", true},

	{"CONNECT", "/", true},
	{"CONNECT", "/index.html", true},
	{"CONNECT", "/icons/ubuntu-logo.png", true},
	{"CONNECT", "/not-existed", true},

	{"INVALID", "/", true},
	{"INVALID", "/index.html", true},
	{"INVALID", "/not-existed", true},
}

func TestMimicApache(t *testing.T) {
	s := proxy.NewServer("anytoken", 100000, 30*time.Second, false, false)
	chListenOn := make(chan string)
	go func() {
		err := s.ServeHTTP(":0", &chListenOn)
		assert.NoError(t, err, "should start chained server")
	}()
	buf := request(t, <-chListenOn)
	ioutil.WriteFile(current, buf.Bytes(), os.ModePerm)
	compare(t, normalize(buf), template)
}

func TestRealApache(t *testing.T) {
	t.Skip("comment out this line and run 'go test -run RealApache' when you want to record http traffic against real apache server")
	addr := "128.199.100.121:80"
	buf := request(t, addr)
	ioutil.WriteFile(target, buf.Bytes(), os.ModePerm)
	ioutil.WriteFile(template, normalize(buf).Bytes(), os.ModePerm)
}

func compare(t *testing.T, buf *bytes.Buffer, file string) {
	c := exec.Command("diff", "-a", "-", file)
	in, err := c.StdinPipe()
	if assert.NoError(t, err, "should gain stdin pipe") {
		go func() {
			buf.WriteTo(in)
			in.Close()
		}()
		b, err := c.CombinedOutput()
		assert.NoError(t, err, "should not different from an real apache")
		t.Log(string(b))
	}
}

func request(t *testing.T, addr string) *bytes.Buffer {
	var buf bytes.Buffer
	for _, c := range candidates {
		conn, err := net.Dial("tcp", addr)
		if assert.NoError(t, err, "should connect") {
			defer conn.Close()
			req := c.method + " " + c.path + " HTTP/1.1\n"
			buf.WriteString(req)
			buf.WriteString("--------------------\n")
			if c.withHostHeader {
				req = req + "Host: " + addr + "\n"
			}
			_, err := conn.Write([]byte(req + "\n"))
			if assert.NoError(t, err, "should write") {
				_, err := io.Copy(&buf, conn)
				assert.NoError(t, err, "should copy")
				buf.WriteString("====================\n\n")
			}
		}
	}
	return &buf
}

func normalize(in *bytes.Buffer) *bytes.Buffer {
	s := in.String()
	exps := []struct {
		regexp *regexp.Regexp
		with   string
	}{
		// Date: ... GMT
		{regexp.MustCompile(`[A-Z][a-z]{2}, \d{2} [A-Z][a-z]{2} \d{4} \d{2}:\d{2}:\d{2} GMT`), "<GMT>"},
		// ETag: "xxxx-xxxxxxxxxxxxx"
		{regexp.MustCompile(`ETag: "\w{3,4}-\w{12,13}"`), `ETag: "<ETAG>"`},
		// Apache/2.4.7 (Ubuntu) Server at 128.199.100.121 Port 80
		{regexp.MustCompile(`Apache\/\d\.\d\.\d+ \(Ubuntu\) Server at \S+ Port \d{2,}`), "Apache/<VERSION> (Ubuntu) Server at <Host> Port <Port>"},
		// Content-Length: 1243
		{regexp.MustCompile(`Content-Length: \d+`), "Content-Length: <...>"},
	}
	for _, e := range exps {
		s = e.regexp.ReplaceAllString(s, e.with)
	}
	return bytes.NewBufferString(s)
}
<commit_msg>fix mimic test<commit_after>package mimic

import (
	"bytes"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"regexp"
	"testing"
	"time"

	"github.com/getlantern/testify/assert"

	proxy "github.com/getlantern/http-proxy"
)

const target = "test-data/apache-2.4.7-ubuntu14.04.raw"
const template = "test-data/apache-2.4.7-ubuntu14.04.tpl"
const current = "test-data/http-proxy.raw"

type entry struct {
	method         string
	path           string
	withHostHeader bool
}

/*
For rationale behind these tests, refer following:
[1] https://github.com/getlantern/lantern/issues/1255#issuecomment-30658217
[2] https://github.com/getlantern/lantern-java/blob/master/src/main/java/org/lantern/proxy/GiveModeHttpFilters.java#L167
[3] https://github.com/apache/httpd/blob/2.4.17/server/core.c#L4379
*/
var candidates = []entry{
	// Without Host header, Apache will return 400
	{"GET", "/", false},
	{"GET", "/index.html", false},

	// 200 with default page
	{"GET", "/", true},
	{"GET", "/index.html", true},

	{"GET", "/icons/ubuntu-logo.png", true},

	// 404
	{"GET", "//cgi-bin/php", true},
	{"GET", "//cgi-bin/php5", true},
	{"GET", "//cgi-bin/php-cgi", true},
	{"GET", "//cgi-bin/php.cgi", true},
	{"GET", "//cgi-bin/php4", true},

	// multiple slashes
	{"GET", "///cgi-bin/php4", true},
	{"GET", "//cgi-bin//php4", true},

	{"GET", "/not-existed", true},
	{"GET", "/end-with-slash/", true},

	{"HEAD", "/", true},
	{"HEAD", "/index.html", true},
	{"HEAD", "/icons/ubuntu-logo.png", true},
	{"HEAD", "/not-existed", true},

	{"POST", "/", true},
	{"POST", "/index.html", true},
	{"POST", "/icons/ubuntu-logo.png", true},
	{"POST", "/not-existed", true},

	{"OPTIONS", "/", true},
	{"OPTIONS", "/index.html", true},
	{"OPTIONS", "/icons/ubuntu-logo.png", true},
	{"OPTIONS", "/not-existed", true},

	{"PUT", "/", true},
	{"PUT", "/index.html", true},
	{"PUT", "/icons/ubuntu-logo.png", true},
	{"PUT", "/not-existed", true},

	{"CONNECT", "/", true},
	{"CONNECT", "/index.html", true},
	{"CONNECT", "/icons/ubuntu-logo.png", true},
	{"CONNECT", "/not-existed", true},

	{"INVALID", "/", true},
	{"INVALID", "/index.html", true},
	{"INVALID", "/not-existed", true},
}

func TestMimicApache(t *testing.T) {
	s := proxy.NewServer("anytoken", 100000, 30*time.Second, true, false)
	chListenOn := make(chan string)
	go func() {
		err := s.ServeHTTP(":0", &chListenOn)
		assert.NoError(t, err, "should start chained server")
	}()
	buf := request(t, <-chListenOn)
	ioutil.WriteFile(current, buf.Bytes(), os.ModePerm)
	compare(t, normalize(buf), template)
}

func TestRealApache(t *testing.T) {
	t.Skip("comment out this line and run 'go test -run RealApache' when you want to record http traffic against real apache server")
	addr := "128.199.100.121:80"
	buf := request(t, addr)
	ioutil.WriteFile(target, buf.Bytes(), os.ModePerm)
	ioutil.WriteFile(template, normalize(buf).Bytes(), os.ModePerm)
}

func compare(t *testing.T, buf *bytes.Buffer, file string) {
	c := exec.Command("diff", "-a", "-", file)
	in, err := c.StdinPipe()
	if assert.NoError(t, err, "should gain stdin pipe") {
		go func() {
			buf.WriteTo(in)
			in.Close()
		}()
		b, err := c.CombinedOutput()
		assert.NoError(t, err, "should not different from an real apache")
		t.Log(string(b))
	}
}

func request(t *testing.T, addr string) *bytes.Buffer {
	var buf bytes.Buffer
	for _, c := range candidates {
		conn, err := net.Dial("tcp", addr)
		if assert.NoError(t, err, "should connect") {
			defer conn.Close()
			req := c.method + " " + c.path + " HTTP/1.1\n"
			buf.WriteString(req)
			buf.WriteString("--------------------\n")
			if c.withHostHeader {
				req = req + "Host: " + addr + "\n"
			}
			_, err := conn.Write([]byte(req + "\n"))
			if assert.NoError(t, err, "should write") {
				_, err := io.Copy(&buf, conn)
				assert.NoError(t, err, "should copy")
				buf.WriteString("====================\n\n")
			}
		}
	}
	return &buf
}

func normalize(in *bytes.Buffer) *bytes.Buffer {
	s := in.String()
	exps := []struct {
		regexp *regexp.Regexp
		with   string
	}{
		// Date: ... GMT
		{regexp.MustCompile(`[A-Z][a-z]{2}, \d{2} [A-Z][a-z]{2} \d{4} \d{2}:\d{2}:\d{2} GMT`), "<GMT>"},
		// ETag: "xxxx-xxxxxxxxxxxxx"
		{regexp.MustCompile(`ETag: "\w{3,4}-\w{12,13}"`), `ETag: "<ETAG>"`},
		// Apache/2.4.7 (Ubuntu) Server at 128.199.100.121 Port 80
		{regexp.MustCompile(`Apache\/\d\.\d\.\d+ \(Ubuntu\) Server at \S+ Port \d{2,}`), "Apache/<VERSION> (Ubuntu) Server at <Host> Port <Port>"},
		// Content-Length: 1243
		{regexp.MustCompile(`Content-Length: \d+`), "Content-Length: <...>"},
	}
	for _, e := range exps {
		s = e.regexp.ReplaceAllString(s, e.with)
	}
	return bytes.NewBufferString(s)
}
<|endoftext|>


<commit_before>package pop

import (
	"fmt"
	"reflect"
	"regexp"
	"strconv"

	"github.com/satori/go.uuid"
)

// Find the first record of the model in the database with a particular id.
//
//	c.Find(&User{}, 1)
func (c *Connection) Find(model interface{}, id interface{}) error {
	return Q(c).Find(model, id)
}

// Find the first record of the model in the database with a particular id.
//
//	q.Find(&User{}, 1)
func (q *Query) Find(model interface{}, id interface{}) error {
	m := &Model{Value: model}
	idq := fmt.Sprintf("%s.id = ?", m.TableName())
	switch t := id.(type) {
	case uuid.UUID:
		return q.Where(idq, t.String()).First(model)
	case string:
		var err error
		id, err = strconv.Atoi(t)
		if err != nil {
			return q.Where(idq, t).First(model)
		}
	}
	return q.Where(idq, id).First(model)
}

// First record of the model in the database that matches the query.
//
//	c.First(&User{})
func (c *Connection) First(model interface{}) error {
	return Q(c).First(model)
}

// First record of the model in the database that matches the query.
//
//	q.Where("name = ?", "mark").First(&User{})
func (q *Query) First(model interface{}) error {
	return q.Connection.timeFunc("First", func() error {
		q.Limit(1)
		m := &Model{Value: model}
		if err := q.Connection.Dialect.SelectOne(q.Connection.Store, m, *q); err != nil {
			return err
		}
		return m.afterFind(q.Connection)
	})
}

// Last record of the model in the database that matches the query.
//
//	c.Last(&User{})
func (c *Connection) Last(model interface{}) error {
	return Q(c).Last(model)
}

// Last record of the model in the database that matches the query.
//
//	q.Where("name = ?", "mark").Last(&User{})
func (q *Query) Last(model interface{}) error {
	return q.Connection.timeFunc("Last", func() error {
		q.Limit(1)
		q.Order("id desc")
		m := &Model{Value: model}
		if err := q.Connection.Dialect.SelectOne(q.Connection.Store, m, *q); err != nil {
			return err
		}
		return m.afterFind(q.Connection)
	})
}

// All retrieves all of the records in the database that match the query.
//
//	c.All(&[]User{})
func (c *Connection) All(models interface{}) error {
	return Q(c).All(models)
}

// All retrieves all of the records in the database that match the query.
//
//	q.Where("name = ?", "mark").All(&[]User{})
func (q *Query) All(models interface{}) error {
	return q.Connection.timeFunc("All", func() error {
		m := &Model{Value: models}
		err := q.Connection.Dialect.SelectMany(q.Connection.Store, m, *q)
		if err == nil && q.Paginator != nil {
			ct, err := q.Count(models)
			if err == nil {
				q.Paginator.TotalEntriesSize = ct
				st := reflect.ValueOf(models).Elem()
				q.Paginator.CurrentEntriesSize = st.Len()
				q.Paginator.TotalPages = (q.Paginator.TotalEntriesSize / q.Paginator.PerPage)
				if q.Paginator.TotalEntriesSize%q.Paginator.PerPage > 0 {
					q.Paginator.TotalPages = q.Paginator.TotalPages + 1
				}
			}
		}
		if err != nil {
			return err
		}
		return m.afterFind(q.Connection)
	})
}

// Exists returns true/false if a record exists in the database that matches
// the query.
//
// 	q.Where("name = ?", "mark").Exists(&User{})
func (q *Query) Exists(model interface{}) (bool, error) {
	i, err := q.Count(model)
	return i != 0, err
}

// Count the number of records in the database.
//
//	c.Count(&User{})
func (c *Connection) Count(model interface{}) (int, error) {
	return Q(c).Count(model)
}

// Count the number of records in the database.
//
//	q.Where("name = ?", "mark").Count(&User{})
func (q Query) Count(model interface{}) (int, error) {
	return q.CountByField(model, "*")
}

func (q Query) CountByField(model interface{}, field string) (int, error) {
	tmpQuery := Q(q.Connection)
	q.Clone(tmpQuery) //avoid mendling with original query

	res := &rowCount{}
	err := tmpQuery.Connection.timeFunc("Count", func() error {
		tmpQuery.Paginator = nil
		tmpQuery.orderClauses = clauses{}
		tmpQuery.limitResults = 0
		query, args := tmpQuery.ToSQL(&Model{Value: model})
		//when query contains custom selected fields / executed using RawQuery,
		//	sql may already contains limit and offset
		r, _ := regexp.Compile("(?i)(limit [0-9]+ offset [0-9]+)$")
		if r.MatchString(query) {
			foundLimit := r.FindString(query)
			query = query[0 : len(query)-len(foundLimit)]
		} else {
			r, _ := regexp.Compile("(?i)(limit [0-9]+)$")
			foundLimit := r.FindString(query)
			query = query[0 : len(query)-len(foundLimit)]
		}

		countQuery := fmt.Sprintf("select count(%s) as row_count from (%s) a", field, query)
		Log(countQuery, args...)
		return q.Connection.Store.Get(res, countQuery, args...)
	})
	return res.Count, err
}

type rowCount struct {
	Count int `db:"row_count"`
}
<commit_msg>outer rlimit n rlimitoffset<commit_after>package pop

import (
	"fmt"
	"reflect"
	"regexp"
	"strconv"

	"github.com/satori/go.uuid"
)

var rLimitOffset = regexp.MustCompile("(?i)(limit [0-9]+ offset [0-9]+)$")
var rLimit = regexp.MustCompile("(?i)(limit [0-9]+)$")

// Find the first record of the model in the database with a particular id.
//
//	c.Find(&User{}, 1)
func (c *Connection) Find(model interface{}, id interface{}) error {
	return Q(c).Find(model, id)
}

// Find the first record of the model in the database with a particular id.
//
//	q.Find(&User{}, 1)
func (q *Query) Find(model interface{}, id interface{}) error {
	m := &Model{Value: model}
	idq := fmt.Sprintf("%s.id = ?", m.TableName())
	switch t := id.(type) {
	case uuid.UUID:
		return q.Where(idq, t.String()).First(model)
	case string:
		var err error
		id, err = strconv.Atoi(t)
		if err != nil {
			return q.Where(idq, t).First(model)
		}
	}
	return q.Where(idq, id).First(model)
}

// First record of the model in the database that matches the query.
//
//	c.First(&User{})
func (c *Connection) First(model interface{}) error {
	return Q(c).First(model)
}

// First record of the model in the database that matches the query.
//
//	q.Where("name = ?", "mark").First(&User{})
func (q *Query) First(model interface{}) error {
	return q.Connection.timeFunc("First", func() error {
		q.Limit(1)
		m := &Model{Value: model}
		if err := q.Connection.Dialect.SelectOne(q.Connection.Store, m, *q); err != nil {
			return err
		}
		return m.afterFind(q.Connection)
	})
}

// Last record of the model in the database that matches the query.
//
//	c.Last(&User{})
func (c *Connection) Last(model interface{}) error {
	return Q(c).Last(model)
}

// Last record of the model in the database that matches the query.
//
//	q.Where("name = ?", "mark").Last(&User{})
func (q *Query) Last(model interface{}) error {
	return q.Connection.timeFunc("Last", func() error {
		q.Limit(1)
		q.Order("id desc")
		m := &Model{Value: model}
		if err := q.Connection.Dialect.SelectOne(q.Connection.Store, m, *q); err != nil {
			return err
		}
		return m.afterFind(q.Connection)
	})
}

// All retrieves all of the records in the database that match the query.
//
//	c.All(&[]User{})
func (c *Connection) All(models interface{}) error {
	return Q(c).All(models)
}

// All retrieves all of the records in the database that match the query.
//
//	q.Where("name = ?", "mark").All(&[]User{})
func (q *Query) All(models interface{}) error {
	return q.Connection.timeFunc("All", func() error {
		m := &Model{Value: models}
		err := q.Connection.Dialect.SelectMany(q.Connection.Store, m, *q)
		if err == nil && q.Paginator != nil {
			ct, err := q.Count(models)
			if err == nil {
				q.Paginator.TotalEntriesSize = ct
				st := reflect.ValueOf(models).Elem()
				q.Paginator.CurrentEntriesSize = st.Len()
				q.Paginator.TotalPages = (q.Paginator.TotalEntriesSize / q.Paginator.PerPage)
				if q.Paginator.TotalEntriesSize%q.Paginator.PerPage > 0 {
					q.Paginator.TotalPages = q.Paginator.TotalPages + 1
				}
			}
		}
		if err != nil {
			return err
		}
		return m.afterFind(q.Connection)
	})
}

// Exists returns true/false if a record exists in the database that matches
// the query.
//
// 	q.Where("name = ?", "mark").Exists(&User{})
func (q *Query) Exists(model interface{}) (bool, error) {
	i, err := q.Count(model)
	return i != 0, err
}

// Count the number of records in the database.
//
//	c.Count(&User{})
func (c *Connection) Count(model interface{}) (int, error) {
	return Q(c).Count(model)
}

// Count the number of records in the database.
//
//	q.Where("name = ?", "mark").Count(&User{})
func (q Query) Count(model interface{}) (int, error) {
	return q.CountByField(model, "*")
}

func (q Query) CountByField(model interface{}, field string) (int, error) {
	tmpQuery := Q(q.Connection)
	q.Clone(tmpQuery) //avoid mendling with original query

	res := &rowCount{}

	err := tmpQuery.Connection.timeFunc("CountByField", func() error {
		tmpQuery.Paginator = nil
		tmpQuery.orderClauses = clauses{}
		tmpQuery.limitResults = 0
		query, args := tmpQuery.ToSQL(&Model{Value: model})
		//when query contains custom selected fields / executed using RawQuery,
		//	sql may already contains limit and offset

		if rLimitOffset.MatchString(query) {
			foundLimit := rLimitOffset.FindString(query)
			query = query[0 : len(query)-len(foundLimit)]
		} else {
			rLimit, _ := regexp.Compile("(?i)(limit [0-9]+)$")
			foundLimit := rLimit.FindString(query)
			query = query[0 : len(query)-len(foundLimit)]
		}

		countQuery := fmt.Sprintf("select count(%s) as row_count from (%s) a", field, query)
		Log(countQuery, args...)
		return q.Connection.Store.Get(res, countQuery, args...)
	})
	return res.Count, err
}

type rowCount struct {
	Count int `db:"row_count"`
}
<|endoftext|>


<commit_before>package hybrid

import (
	"crypto/rsa"
	"encoding/asn1"
	"github.com/kisom/gocrypto/chapter7/pkc"
	"github.com/kisom/gocrypto/chapter8/pks"
)

const (
	SymKeyLen = 16 // AES-128 key size
	MacKeyLen = 32 // HMAC-SHA256 key size
)

type Message struct {
	Key []byte
	Sig []byte
	Msg []byte
}

func generateEphemeralKeys(prv *rsa.PrivateKey) (key, sig []byte, err error) {
        var sym, mac []byte
	if sym, err = Random(SymKeyLen); err != nil {
		return
	} else if mac, err = Random(MacKeyLen); err != nil {
		return
	}

	key = make([]byte, SymKeyLen+MacKeyLen)
	copy(key, sym)
	copy(key[SymKeyLen:], mac)
	sig, err = pks.Sign(prv, key)
	return
}

func Encrypt(prv *rsa.PrivateKey, pub *rsa.PublicKey, m []byte) (ct []byte, err error) {
	var msg Message
	var key []byte
	if key, msg.Sig, err = generateEphemeralKeys(prv); err != nil {
		return
	}

	if msg.Key, err = pkc.Encrypt(pub, key); err != nil {
		return
	} else if msg.Msg, err = symEncrypt(key[:SymKeyLen], key[SymKeyLen:], m); err != nil {
		return
	}
	ct, err = asn1.Marshal(msg)
	scrub(key, 3)
	return
}

func readEphemeralKeys(prv *rsa.PrivateKey, pub *rsa.PublicKey, key, sig []byte) (sym, mac []byte, err error) {
	if key, err = pkc.Decrypt(prv, key); err != nil {
		return
	}

	err = pks.Verify(pub, key, sig)
	if err != nil {
		return
	}
	sym = key[:SymKeyLen]
	mac = key[SymKeyLen:]
	return
}

func Decrypt(prv *rsa.PrivateKey, pub *rsa.PublicKey, ct []byte) (m []byte, err error) {
	var msg Message

	if _, err = asn1.Unmarshal(ct, &msg); err != nil {
		return
	}

	var sym, mac []byte
	sym, mac, err = readEphemeralKeys(prv, pub, msg.Key, msg.Sig)
	if err != nil {
		return
	}
	m, err = symDecrypt(sym, mac, msg.Msg)
	scrub(sym, 3)
	scrub(mac, 3)
	return
}
<commit_msg>Go fmt.<commit_after>package hybrid

import (
	"crypto/rsa"
	"encoding/asn1"
	"github.com/kisom/gocrypto/chapter7/pkc"
	"github.com/kisom/gocrypto/chapter8/pks"
)

const (
	SymKeyLen = 16 // AES-128 key size
	MacKeyLen = 32 // HMAC-SHA256 key size
)

type Message struct {
	Key []byte
	Sig []byte
	Msg []byte
}

func generateEphemeralKeys(prv *rsa.PrivateKey) (key, sig []byte, err error) {
	var sym, mac []byte
	if sym, err = Random(SymKeyLen); err != nil {
		return
	} else if mac, err = Random(MacKeyLen); err != nil {
		return
	}

	key = make([]byte, SymKeyLen+MacKeyLen)
	copy(key, sym)
	copy(key[SymKeyLen:], mac)
	sig, err = pks.Sign(prv, key)
	return
}

func Encrypt(prv *rsa.PrivateKey, pub *rsa.PublicKey, m []byte) (ct []byte, err error) {
	var msg Message
	var key []byte
	if key, msg.Sig, err = generateEphemeralKeys(prv); err != nil {
		return
	}

	if msg.Key, err = pkc.Encrypt(pub, key); err != nil {
		return
	} else if msg.Msg, err = symEncrypt(key[:SymKeyLen], key[SymKeyLen:], m); err != nil {
		return
	}
	ct, err = asn1.Marshal(msg)
	scrub(key, 3)
	return
}

func readEphemeralKeys(prv *rsa.PrivateKey, pub *rsa.PublicKey, key, sig []byte) (sym, mac []byte, err error) {
	if key, err = pkc.Decrypt(prv, key); err != nil {
		return
	}

	err = pks.Verify(pub, key, sig)
	if err != nil {
		return
	}
	sym = key[:SymKeyLen]
	mac = key[SymKeyLen:]
	return
}

func Decrypt(prv *rsa.PrivateKey, pub *rsa.PublicKey, ct []byte) (m []byte, err error) {
	var msg Message

	if _, err = asn1.Unmarshal(ct, &msg); err != nil {
		return
	}

	var sym, mac []byte
	sym, mac, err = readEphemeralKeys(prv, pub, msg.Key, msg.Sig)
	if err != nil {
		return
	}
	m, err = symDecrypt(sym, mac, msg.Msg)
	scrub(sym, 3)
	scrub(mac, 3)
	return
}
<|endoftext|>


<commit_before><commit_msg>Add support for rkt export<commit_after><|endoftext|>


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

package service

import (
	"testing"
	"time"

	"github.com/control-center/serviced/coordinator/client"
	"github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/domain/host"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/domain/servicestate"
)

func TestGetServiceStatus(t *testing.T) {
	conn := client.NewTestConnection()
	defer conn.Close()

	// Create the service
	svc := &service.Service{
		ID:        "test-service-1",
		Endpoints: make([]service.ServiceEndpoint, 1),
	}

	if err := UpdateService(conn, svc); err != nil {
		t.Fatalf("Could not add service %s: %s", svc.ID, err)
	}
	if err := AddHost(conn, &host.Host{ID: "test-host-1"}); err != nil {
		t.Fatalf("Could not register host: %s", err)
	}

	// 0 states
	statusmap, err := GetServiceStatus(conn, svc.ID)
	if err != nil {
		t.Fatalf("Could not get status for service %s: %s", svc.ID, err)
	}

	if statusmap != nil && len(statusmap) > 0 {
		t.Errorf("Expected 0 statuses returned; got %d", len(statusmap))
	}

	// Add service states
	var states []*servicestate.ServiceState
	for i := 0; i < 3; i++ {
		state, err := servicestate.BuildFromService(svc, "test-host-1")
		if err != nil {
			t.Fatalf("Could not generate instance from service %s", svc.ID)
		} else if err := addInstance(conn, state); err != nil {
			t.Fatalf("Could not add instance %s from service %s", state.ID, state.ServiceID)
		}
		states = append(states, state)
	}

	expected := make(map[string]dao.Status)
	// State 0 started
	states[0].Started = time.Now()
	if err := UpdateServiceState(conn, states[0]); err != nil {
		t.Fatalf("Could not \"start\" service state %s: %s", states[0].ID, err)
	}
	expected[states[0].ID] = dao.Running

	// State 1 paused
	states[1].Started = time.Now()
	states[1].Paused = true
	if err := UpdateServiceState(conn, states[1]); err != nil {
		t.Fatalf("Could not \"pause\" service state %s: %s", states[1].ID, err)
	}
	expected[states[1].ID] = dao.Resuming

	// State 2 stopped (no update)
	expected[states[2].ID] = dao.Starting

	t.Log("Desired state is RUN")
	statusmap, err = GetServiceStatus(conn, svc.ID)
	if err != nil {
		t.Fatalf("Could not get the status for service %s: %s", svc.ID, err)
	} else if len(statusmap) != len(states) {
		t.Errorf("MISMATCH: expected %d states; actual %d", len(states), len(statusmap))
	}

	// Verify
	for _, svcstatus := range statusmap {
		expect, ok := expected[svcstatus.State.ID]
		if !ok {
			t.Fatalf("Missing service state %s", svcstatus.State.ID)
		} else if expect != svcstatus.Status {
			t.Errorf("MISMATCH: expected %s; actual %s", expect, svcstatus.Status)
		}
	}

	t.Log("Desired state is PAUSE")
	for _, state := range states {
		if err := pauseInstance(conn, state.HostID, state.ID); err != nil {
			t.Fatalf("Could not pause instance %s: %s", state.ID, err)
		}
	}
	expected[states[0].ID] = dao.Pausing
	expected[states[1].ID] = dao.Paused
	expected[states[2].ID] = dao.Stopped

	statusmap, err = GetServiceStatus(conn, svc.ID)
	if err != nil {
		t.Fatalf("Could not get the status for service %s: %s", svc.ID, err)
	} else if len(statusmap) != len(states) {
		t.Errorf("MISMATCH: expected %d states; actual %d", len(states), len(statusmap))
	}

	// Verify
	for _, svcstatus := range statusmap {
		expect, ok := expected[svcstatus.State.ID]
		if !ok {
			t.Fatalf("Missing service state %s", svcstatus.State.ID)
		} else if expect != svcstatus.Status {
			t.Errorf("MISMATCH: expected %s; actual %s", expect, svcstatus.Status)
		}
	}

	t.Log("Desired state is STOP")
	for _, state := range states {
		if err := StopServiceInstance(conn, state.HostID, state.ID); err != nil {
			t.Fatalf("Could not stop instance %s: %s", state.ID, err)
		}
	}
	expected[states[0].ID] = dao.Stopping
	expected[states[1].ID] = dao.Stopping
	expected[states[2].ID] = dao.Stopped

	statusmap, err = GetServiceStatus(conn, svc.ID)
	if err != nil {
		t.Fatalf("Could not get the status for service %s: %s", svc.ID, err)
	} else if len(statusmap) != len(states) {
		t.Errorf("MISMATCH: expected %d states; actual %d", len(states), len(statusmap))
	}

	// Verify
	for _, svcstatus := range statusmap {
		expect, ok := expected[svcstatus.State.ID]
		if !ok {
			t.Fatalf("Missing service state %s", svcstatus.State.ID)
		} else if expect != svcstatus.Status {
			t.Errorf("MISMATCH: expected %s; actual %s", expect, svcstatus.Status)
		}
	}

}
<commit_msg>Updated unit tests for servicestate in zzk/service<commit_after>// Copyright 2014 The Serviced Authors.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package service

import (
	"path"
	"time"

	"github.com/control-center/serviced/dao"
	"github.com/control-center/serviced/domain/host"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/domain/servicestate"
	"github.com/control-center/serviced/zzk"

	. "gopkg.in/check.v1"
)

func (t *ZZKTest) TestGetServiceStatus(c *C) {
	conn, err := zzk.GetLocalConnection("/")
	c.Assert(err, IsNil)

	// Add a service
	svc := service.Service{ID: "test-service-1", Instances: 3}
	err = UpdateService(conn, &svc)
	c.Assert(err, IsNil)

	// Add a host
	register := func(hostID string) string {
		c.Logf("Registering host %s", hostID)
		host := host.Host{ID: hostID}
		err := AddHost(conn, &host)
		c.Assert(err, IsNil)
		p, err := conn.CreateEphemeral(hostregpath(hostID), &HostNode{Host: &host})
		c.Assert(err, IsNil)
		return path.Base(p)
	}
	register("test-host-1")

	// Case 1: Zero service states
	statusmap, err := GetServiceStatus(conn, svc.ID)
	c.Assert(err, IsNil)
	c.Assert(statusmap, NotNil)
	c.Assert(statusmap, DeepEquals, map[string]dao.ServiceStatus{})

	// Add states
	addstates := func(hostID string, svc *service.Service, count int) []string {
		c.Logf("Adding %d service states for service %s on host %s", count, svc.ID, hostID)
		stateIDs := make([]string, count)
		for i := 0; i < count; i++ {
			state, err := servicestate.BuildFromService(svc, hostID)
			c.Assert(err, IsNil)
			err = addInstance(conn, state)
			c.Assert(err, IsNil)
			_, err = LoadRunningService(conn, state.ServiceID, state.ID)
			c.Assert(err, IsNil)
			stateIDs[i] = state.ID
		}
		return stateIDs
	}
	stateIDs := addstates("test-host-1", &svc, 3)

	states := make(map[string]servicestate.ServiceState)
	updatestate := func(stateID, serviceID string, update func(state *servicestate.ServiceState)) servicestate.ServiceState {
		var node ServiceStateNode
		err = conn.Get(servicepath(serviceID, stateID), &node)
		c.Assert(err, IsNil)
		update(node.ServiceState)
		err = conn.Set(servicepath(serviceID, stateID), &node)
		c.Assert(err, IsNil)
		return *node.ServiceState
	}

	// State 0 started
	states[stateIDs[0]] = updatestate(stateIDs[0], svc.ID, func(s *servicestate.ServiceState) {
		s.Started = time.Now()
	})

	// State 1 paused
	states[stateIDs[1]] = updatestate(stateIDs[1], svc.ID, func(s *servicestate.ServiceState) {
		s.Started = time.Now()
		s.Paused = true
	})

	// State 2 stopped (no update)
	states[stateIDs[2]] = updatestate(stateIDs[2], svc.ID, func(s *servicestate.ServiceState) {})

	c.Log("Desired state is RUN")
	statusmap, err = GetServiceStatus(conn, svc.ID)
	c.Assert(err, IsNil)
	c.Assert(statusmap, DeepEquals, map[string]dao.ServiceStatus{
		stateIDs[0]: dao.ServiceStatus{states[stateIDs[0]], dao.Running},
		stateIDs[1]: dao.ServiceStatus{states[stateIDs[1]], dao.Resuming},
		stateIDs[2]: dao.ServiceStatus{states[stateIDs[2]], dao.Starting},
	})

	c.Log("Desired state is PAUSE")
	for _, state := range stateIDs {
		err := pauseInstance(conn, "test-host-1", state)
		c.Assert(err, IsNil)
	}
	statusmap, err = GetServiceStatus(conn, svc.ID)
	c.Assert(err, IsNil)
	c.Assert(statusmap, DeepEquals, map[string]dao.ServiceStatus{
		stateIDs[0]: dao.ServiceStatus{states[stateIDs[0]], dao.Pausing},
		stateIDs[1]: dao.ServiceStatus{states[stateIDs[1]], dao.Paused},
		stateIDs[2]: dao.ServiceStatus{states[stateIDs[2]], dao.Stopped},
	})

	c.Log("Desired state is STOP")
	for _, state := range stateIDs {
		err := StopServiceInstance(conn, "test-host-1", state)
		c.Assert(err, IsNil)
	}
	statusmap, err = GetServiceStatus(conn, svc.ID)
	c.Assert(err, IsNil)
	c.Assert(statusmap, DeepEquals, map[string]dao.ServiceStatus{
		stateIDs[0]: dao.ServiceStatus{states[stateIDs[0]], dao.Stopping},
		stateIDs[1]: dao.ServiceStatus{states[stateIDs[1]], dao.Stopping},
		stateIDs[2]: dao.ServiceStatus{states[stateIDs[2]], dao.Stopped},
	})
}
<|endoftext|>


<commit_before>// bulk_load_influx loads an InfluxDB daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"
)

// Program option vars:
var (
	daemonUrl string
	dbName    string
	workers   int
	batchSize int
	backoff   time.Duration
	doLoad    bool
)

// Global vars
var (
	bufPool        sync.Pool
	batchChan      chan *bytes.Buffer
	inputDone      chan struct{}
	workersGroup   sync.WaitGroup
	backingOffChan chan bool
	backingOffDone chan struct{}
)

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "http://localhost:8086", "Influxd URL.")
	flag.StringVar(&dbName, "db", "benchmark_db", "Database name.")
	flag.IntVar(&batchSize, "batch-size", 5000, "Batch size (input lines).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.DurationVar(&backoff, "backoff", time.Second, "Time to sleep between requests when server indicates backpressure is needed.")
	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.Parse()
}

func main() {
	if doLoad {
		// check that there are no pre-existing databases:
		existingDatabases, err := listDatabases(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		if len(existingDatabases) > 0 {
			log.Fatalf("There are databases already in the data store. If you know what you are doing, run the command:\ncurl 'http://localhost:8086/query?q=drop%%20database%%20%s'\n", existingDatabases[0])
		}

		err = createDb(daemonUrl, dbName)
		if err != nil {
			log.Fatal(err)
		}
	}

	bufPool = sync.Pool{
		New: func() interface{} {
			return bytes.NewBuffer(make([]byte, 0, 4*1024*1024))
		},
	}

	batchChan = make(chan *bytes.Buffer, workers)
	inputDone = make(chan struct{})

	backingOffChan = make(chan bool, 100)
	backingOffDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		cfg := HTTPWriterConfig{
			Host:     daemonUrl,
			Database: dbName,
		}
		go processBatches(NewHTTPWriter(cfg))
	}

	go processBackoffMessages()

	start := time.Now()
	itemsRead := scan(batchSize)

	<-inputDone
	close(batchChan)

	workersGroup.Wait()

	close(backingOffChan)
	<-backingOffDone

	end := time.Now()
	took := end.Sub(start)
	rate := float64(itemsRead) / float64(took.Seconds())

	fmt.Printf("loaded %d items in %fsec with %d workers (mean rate %f/sec)\n", itemsRead, took.Seconds(), workers, rate)
}

// scan reads one line at a time from stdin.
// When the requested number of lines per batch is met, send a batch over batchChan for the workers to write.
func scan(linesPerBatch int) int64 {
	buf := bufPool.Get().(*bytes.Buffer)

	var n int
	var itemsRead int64
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		itemsRead++

		buf.Write(scanner.Bytes())
		buf.Write([]byte("\n"))

		n++
		if n >= linesPerBatch {
			batchChan <- buf
			buf = bufPool.Get().(*bytes.Buffer)
			n = 0
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading input: %s", err.Error())
	}

	// Finished reading input, make sure last batch goes out.
	if n > 0 {
		batchChan <- buf
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)

	return itemsRead
}

// processBatches reads byte buffers from batchChan and writes them to the target server, while tracking stats on the write.
func processBatches(w LineProtocolWriter) {
	for batch := range batchChan {
		if !doLoad {
			continue
		}
		// Write the batch: try until backoff is not needed.
		var err error
		for {
			_, err = w.WriteLineProtocol(batch.Bytes())
			if err == BackoffError {
				backingOffChan <- true
				time.Sleep(backoff)
			} else {
				backingOffChan <- false
				break
			}
		}
		if err != nil {
			log.Fatalf("Error writing: %s\n", err.Error())
		}

		// Return the batch buffer to the pool.
		batch.Reset()
		bufPool.Put(batch)
	}
	workersGroup.Done()
}

func processBackoffMessages() {
	var start time.Time
	last := false
	for this := range backingOffChan {
		if this && !last {
			fmt.Println("backoff started")
			start = time.Now()
			last = true
		} else if !this && last {
			took := time.Now().Sub(start)
			fmt.Printf("backoff ended after %.02fsec\n", took.Seconds())
			start = time.Now()
			last = false
		}
	}
	backingOffDone<-struct{}{}
}

func createDb(daemon_url, dbname string) error {
	u, err := url.Parse(daemon_url)
	if err != nil {
		return err
	}

	// serialize params the right way:
	u.Path = "query"
	v := u.Query()
	v.Set("q", fmt.Sprintf("CREATE DATABASE %s", dbname))
	u.RawQuery = v.Encode()

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	// does the body need to be read into the void?

	if resp.StatusCode != 200 {
		return fmt.Errorf("bad db create")
	}
	return nil
}

// listDatabases lists the existing databases in InfluxDB.
func listDatabases(daemonUrl string) ([]string, error) {
	u := fmt.Sprintf("%s/query?q=show%%20databases", daemonUrl)
	resp, err := http.Get(u)
	if err != nil {
		return nil, fmt.Errorf("listDatabases error: %s", err.Error())
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// Do ad-hoc parsing to find existing database names:
	// {"results":[{"series":[{"name":"databases","columns":["name"],"values":[["_internal"],["benchmark_db"]]}]}]}%
	type listingType struct {
		Results []struct {
			Series []struct {
				Values [][]string
			}
		}
	}
	var listing listingType
	err = json.Unmarshal(body, &listing)
	if err != nil {
		return nil, err
	}

	ret := []string{}
	for _, nestedName := range listing.Results[0].Series[0].Values {
		name := nestedName[0]
		// the _internal database is skipped:
		if name == "_internal" {
			continue
		}
		ret = append(ret, name)
	}
	return ret, nil
}
<commit_msg>Improve backoff messages.<commit_after>// bulk_load_influx loads an InfluxDB daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"sync"
	"time"
)

// Program option vars:
var (
	daemonUrl string
	dbName    string
	workers   int
	batchSize int
	backoff   time.Duration
	doLoad    bool
)

// Global vars
var (
	bufPool        sync.Pool
	batchChan      chan *bytes.Buffer
	inputDone      chan struct{}
	workersGroup   sync.WaitGroup
	backingOffChan chan bool
	backingOffDone chan struct{}
)

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "http://localhost:8086", "Influxd URL.")
	flag.StringVar(&dbName, "db", "benchmark_db", "Database name.")
	flag.IntVar(&batchSize, "batch-size", 5000, "Batch size (input lines).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")
	flag.DurationVar(&backoff, "backoff", time.Second, "Time to sleep between requests when server indicates backpressure is needed.")
	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.Parse()
}

func main() {
	if doLoad {
		// check that there are no pre-existing databases:
		existingDatabases, err := listDatabases(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		if len(existingDatabases) > 0 {
			log.Fatalf("There are databases already in the data store. If you know what you are doing, run the command:\ncurl 'http://localhost:8086/query?q=drop%%20database%%20%s'\n", existingDatabases[0])
		}

		err = createDb(daemonUrl, dbName)
		if err != nil {
			log.Fatal(err)
		}
	}

	bufPool = sync.Pool{
		New: func() interface{} {
			return bytes.NewBuffer(make([]byte, 0, 4*1024*1024))
		},
	}

	batchChan = make(chan *bytes.Buffer, workers)
	inputDone = make(chan struct{})

	backingOffChan = make(chan bool, 100)
	backingOffDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		cfg := HTTPWriterConfig{
			Host:     daemonUrl,
			Database: dbName,
		}
		go processBatches(NewHTTPWriter(cfg))
	}

	go processBackoffMessages()

	start := time.Now()
	itemsRead := scan(batchSize)

	<-inputDone
	close(batchChan)

	workersGroup.Wait()

	close(backingOffChan)
	<-backingOffDone

	end := time.Now()
	took := end.Sub(start)
	rate := float64(itemsRead) / float64(took.Seconds())

	fmt.Printf("loaded %d items in %fsec with %d workers (mean rate %f/sec)\n", itemsRead, took.Seconds(), workers, rate)
}

// scan reads one line at a time from stdin.
// When the requested number of lines per batch is met, send a batch over batchChan for the workers to write.
func scan(linesPerBatch int) int64 {
	buf := bufPool.Get().(*bytes.Buffer)

	var n int
	var itemsRead int64
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		itemsRead++

		buf.Write(scanner.Bytes())
		buf.Write([]byte("\n"))

		n++
		if n >= linesPerBatch {
			batchChan <- buf
			buf = bufPool.Get().(*bytes.Buffer)
			n = 0
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading input: %s", err.Error())
	}

	// Finished reading input, make sure last batch goes out.
	if n > 0 {
		batchChan <- buf
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)

	return itemsRead
}

// processBatches reads byte buffers from batchChan and writes them to the target server, while tracking stats on the write.
func processBatches(w LineProtocolWriter) {
	for batch := range batchChan {
		if !doLoad {
			continue
		}
		// Write the batch: try until backoff is not needed.
		var err error
		for {
			_, err = w.WriteLineProtocol(batch.Bytes())
			if err == BackoffError {
				backingOffChan <- true
				time.Sleep(backoff)
			} else {
				backingOffChan <- false
				break
			}
		}
		if err != nil {
			log.Fatalf("Error writing: %s\n", err.Error())
		}

		// Return the batch buffer to the pool.
		batch.Reset()
		bufPool.Put(batch)
	}
	workersGroup.Done()
}

func processBackoffMessages() {
	var totalBackoffSecs float64
	var start time.Time
	last := false
	for this := range backingOffChan {
		if this && !last {
			start = time.Now()
			last = true
		} else if !this && last {
			took := time.Now().Sub(start)
			fmt.Printf("backoff took %.02fsec\n", took.Seconds())
			totalBackoffSecs += took.Seconds()
			last = false
			start = time.Now()
		}
	}
	fmt.Printf("backoffs took a total of %fsec of runtime\n", totalBackoffSecs)
	backingOffDone<-struct{}{}
}

func createDb(daemon_url, dbname string) error {
	u, err := url.Parse(daemon_url)
	if err != nil {
		return err
	}

	// serialize params the right way:
	u.Path = "query"
	v := u.Query()
	v.Set("q", fmt.Sprintf("CREATE DATABASE %s", dbname))
	u.RawQuery = v.Encode()

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	// does the body need to be read into the void?

	if resp.StatusCode != 200 {
		return fmt.Errorf("bad db create")
	}
	return nil
}

// listDatabases lists the existing databases in InfluxDB.
func listDatabases(daemonUrl string) ([]string, error) {
	u := fmt.Sprintf("%s/query?q=show%%20databases", daemonUrl)
	resp, err := http.Get(u)
	if err != nil {
		return nil, fmt.Errorf("listDatabases error: %s", err.Error())
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// Do ad-hoc parsing to find existing database names:
	// {"results":[{"series":[{"name":"databases","columns":["name"],"values":[["_internal"],["benchmark_db"]]}]}]}%
	type listingType struct {
		Results []struct {
			Series []struct {
				Values [][]string
			}
		}
	}
	var listing listingType
	err = json.Unmarshal(body, &listing)
	if err != nil {
		return nil, err
	}

	ret := []string{}
	for _, nestedName := range listing.Results[0].Series[0].Values {
		name := nestedName[0]
		// the _internal database is skipped:
		if name == "_internal" {
			continue
		}
		ret = append(ret, name)
	}
	return ret, nil
}
<|endoftext|>


<commit_before>package main

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

	"github.com/google/go-jsonnet"
	"github.com/google/go-jsonnet/ast"

	// #cgo CXXFLAGS: -std=c++11 -Wall -I../cpp-jsonnet/include
	// #include "internal.h"
	"C"
)

type vm struct {
	*jsonnet.VM
	importer *jsonnet.FileImporter
}

type jsonValue struct {
	val interface{}
	// these objects are exclusively owned by this jsonValue
	owned []*C.struct_JsonnetJsonValue
}

var handles = handlesTable{}
var versionString *C.char

//export jsonnet_version
func jsonnet_version() *C.char {
	if versionString == nil {
		version := jsonnet.Version() + " (go-jsonnet)"
		versionString = C.CString(version)
	}
	return versionString
}

//export jsonnet_make
func jsonnet_make() *C.struct_JsonnetVm {
	newVM := &vm{jsonnet.MakeVM(), &jsonnet.FileImporter{}}
	newVM.Importer(newVM.importer)

	id, err := handles.make(newVM)

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	return C.jsonnet_internal_make_vm_with_id(C.uint32_t(id))
}

//export jsonnet_destroy
func jsonnet_destroy(vmRef *C.struct_JsonnetVm) {
	if err := handles.free(uint32(vmRef.id)); err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	C.jsonnet_internal_free_vm(vmRef)
}

func getVM(vmRef *C.struct_JsonnetVm) *vm {
	ref, err := handles.get(uint32(vmRef.id))

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	v, ok := ref.(*vm)

	if !ok {
		fmt.Fprintf(os.Stderr, "provided handle has a different type")
		os.Exit(1)
	}

	return v
}

func evaluateSnippet(vmRef *C.struct_JsonnetVm, filename string, code string, e *C.int) *C.char {
	vm := getVM(vmRef)
	out, err := vm.EvaluateSnippet(filename, code)
	var result *C.char
	if err != nil {
		*e = 1
		result = C.CString(err.Error())
	} else {
		*e = 0
		result = C.CString(out)
	}
	return result
}

//export jsonnet_evaluate_snippet
func jsonnet_evaluate_snippet(vmRef *C.struct_JsonnetVm, filename *C.char, code *C.char, e *C.int) *C.char {
	f := C.GoString(filename)
	s := C.GoString(code)
	return evaluateSnippet(vmRef, f, s, e)
}

//export jsonnet_evaluate_file
func jsonnet_evaluate_file(vmRef *C.struct_JsonnetVm, filename *C.char, e *C.int) *C.char {
	f := C.GoString(filename)
	data, err := ioutil.ReadFile(f)
	if err != nil {
		*e = 1
		// TODO(sbarzowski) make sure that it's ok allocation-wise
		return C.CString(fmt.Sprintf("Failed to read input file: %s: %s", f, err.Error()))
	}
	return evaluateSnippet(vmRef, f, string(data), e)
}

//export jsonnet_max_stack
func jsonnet_max_stack(vmRef *C.struct_JsonnetVm, v C.uint) {
	vm := getVM(vmRef)
	vm.MaxStack = int(v) // potentially dangerous conversion
}

//export jsonnet_string_output
func jsonnet_string_output(vmRef *C.struct_JsonnetVm, v C.int) {
	vm := getVM(vmRef)
	vm.StringOutput = v != 0
}

//export jsonnet_max_trace
func jsonnet_max_trace(vmRef *C.struct_JsonnetVm, v C.uint) {
	vm := getVM(vmRef)
	vm.ErrorFormatter.SetMaxStackTraceSize(int(v)) // potentially dangerous conversion
}

//export jsonnet_jpath_add
func jsonnet_jpath_add(vmRef *C.struct_JsonnetVm, path *C.char) {
	vm := getVM(vmRef)
	vm.importer.JPaths = append(vm.importer.JPaths, C.GoString(path))
}

//export jsonnet_ext_var
func jsonnet_ext_var(vmRef *C.struct_JsonnetVm, key, value *C.char) {
	vm := getVM(vmRef)
	vm.ExtVar(C.GoString(key), C.GoString(value))
}

//export jsonnet_ext_code
func jsonnet_ext_code(vmRef *C.struct_JsonnetVm, key, value *C.char) {
	vm := getVM(vmRef)
	vm.ExtCode(C.GoString(key), C.GoString(value))
}

//export jsonnet_tla_var
func jsonnet_tla_var(vmRef *C.struct_JsonnetVm, key, value *C.char) {
	vm := getVM(vmRef)
	vm.TLAVar(C.GoString(key), C.GoString(value))
}

//export jsonnet_tla_code
func jsonnet_tla_code(vmRef *C.struct_JsonnetVm, key, value *C.char) {
	vm := getVM(vmRef)
	vm.TLACode(C.GoString(key), C.GoString(value))
}

//export jsonnet_native_callback
func jsonnet_native_callback(vmRef *C.struct_JsonnetVm, name *C.char, cb *C.JsonnetNativeCallback, ctx unsafe.Pointer, params **C.char) {
	vm := getVM(vmRef)
	p := unsafe.Pointer(params)
	sz := unsafe.Sizeof(*params)

	var paramNames ast.Identifiers

	for i := 0; ; i++ {
		param := (**C.char)(unsafe.Pointer(uintptr(p) + uintptr(i)*sz))

		if *param == nil {
			break
		}

		paramNames = append(paramNames, ast.Identifier(C.GoString(*param)))
	}

	f := &jsonnet.NativeFunction{
		Name:   C.GoString(name),
		Params: paramNames,
		Func: func(x []interface{}) (interface{}, error) {
			var (
				arr     []*C.struct_JsonnetJsonValue
				argv    **C.struct_JsonnetJsonValue
				success = C.int(0)
			)

			if len(x) > 0 {
				arr = make([]*C.struct_JsonnetJsonValue, 0, len(x))

				for _, json := range x {
					arr = append(arr, createJSONValue(vmRef, json))
				}

				argv = &(arr[0])
			}

			result := C.jsonnet_internal_execute_native(cb, ctx, argv, &success)
			v := getJSONValue(result)

			for _, val := range arr {
				jsonnet_json_destroy(vmRef, val)
			}

			jsonnet_json_destroy(vmRef, result)

			if success != 1 {
				return nil, fmt.Errorf("failed to execute native callback, code: %d", success)
			}

			return v.val, nil
		},
	}

	vm.NativeFunction(f)
}

//export jsonnet_json_extract_string
func jsonnet_json_extract_string(vmRef *C.struct_JsonnetVm, json *C.struct_JsonnetJsonValue) *C.char {
	v := getJSONValue(json)
	str, ok := v.val.(string)

	if !ok {
		return nil
	}

	return C.CString(str)
}

//export jsonnet_json_extract_number
func jsonnet_json_extract_number(vmRef *C.struct_JsonnetVm, json *C.struct_JsonnetJsonValue, out *C.double) C.int {
	v := getJSONValue(json)

	switch f := v.val.(type) {
	case float64:
		*out = C.double(f)
		return 1
	case int:
		*out = C.double(f)
		return 1
	}

	return 0
}

//export jsonnet_json_extract_bool
func jsonnet_json_extract_bool(vmRef *C.struct_JsonnetVm, json *C.struct_JsonnetJsonValue) C.int {
	v := getJSONValue(json)
	b, ok := v.val.(bool)

	if !ok {
		return 2
	}

	if b {
		return 1
	}

	return 0
}

//export jsonnet_json_extract_null
func jsonnet_json_extract_null(vmRef *C.struct_JsonnetVm, json *C.struct_JsonnetJsonValue) C.int {
	v := getJSONValue(json)

	if v.val == nil {
		return 1
	}

	return 0
}

//export jsonnet_json_make_string
func jsonnet_json_make_string(vmRef *C.struct_JsonnetVm, v *C.char) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, C.GoString(v))
}

//export jsonnet_json_make_number
func jsonnet_json_make_number(vmRef *C.struct_JsonnetVm, v C.double) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, float64(v))
}

//export jsonnet_json_make_bool
func jsonnet_json_make_bool(vmRef *C.struct_JsonnetVm, v C.int) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, v != 0)
}

//export jsonnet_json_make_null
func jsonnet_json_make_null(vmRef *C.struct_JsonnetVm) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, nil)
}

//export jsonnet_json_make_array
func jsonnet_json_make_array(vmRef *C.struct_JsonnetVm) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, []interface{}{})
}

//export jsonnet_json_array_append
func jsonnet_json_array_append(vmRef *C.struct_JsonnetVm, arr *C.struct_JsonnetJsonValue, v *C.struct_JsonnetJsonValue) {
	_ = getVM(vmRef) // Here we check it, is it ok?
	json := getJSONValue(arr)
	slice, ok := json.val.([]interface{})

	if !ok {
		fmt.Fprintf(os.Stderr, "array should be provided")
		os.Exit(1)
	}

	json.val = append(slice, getJSONValue(v).val)
	json.owned = append(json.owned, v)
}

//export jsonnet_json_make_object
func jsonnet_json_make_object(vmRef *C.struct_JsonnetVm) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, make(map[string]interface{}))
}

//export jsonnet_json_object_append
func jsonnet_json_object_append(
	vmRef *C.struct_JsonnetVm,
	obj *C.struct_JsonnetJsonValue,
	f *C.char,
	v *C.struct_JsonnetJsonValue,
) {
	_ = getVM(vmRef) // Here we check it, is it ok?
	d := getJSONValue(obj)
	table, ok := d.val.(map[string]interface{})

	if !ok {
		fmt.Fprintf(os.Stderr, "object should be provided")
		os.Exit(1)
	}

	table[C.GoString(f)] = getJSONValue(v).val
	d.owned = append(d.owned, v)
}

//export jsonnet_json_destroy
func jsonnet_json_destroy(vmRef *C.struct_JsonnetVm, v *C.struct_JsonnetJsonValue) {
	for _, child := range getJSONValue(v).owned {
		jsonnet_json_destroy(vmRef, child)
	}

	if err := handles.free(uint32(v.id)); err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	C.jsonnet_internal_free_json(v)
}

func createJSONValue(vmRef *C.struct_JsonnetVm, val interface{}) *C.struct_JsonnetJsonValue {
	id, err := handles.make(&jsonValue{val: val})

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	return C.jsonnet_internal_make_json_with_id(C.uint32_t(id))
}

func getJSONValue(jsonRef *C.struct_JsonnetJsonValue) *jsonValue {
	ref, err := handles.get(uint32(jsonRef.id))

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	v, ok := ref.(*jsonValue)

	if !ok {
		fmt.Fprintf(os.Stderr, "provided handle has a different type")
		os.Exit(1)
	}

	return v
}

func main() {
}
<commit_msg>Code review fixes: Removed unnecessary vm handle checks<commit_after>package main

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

	"github.com/google/go-jsonnet"
	"github.com/google/go-jsonnet/ast"

	// #cgo CXXFLAGS: -std=c++11 -Wall -I../cpp-jsonnet/include
	// #include "internal.h"
	"C"
)

type vm struct {
	*jsonnet.VM
	importer *jsonnet.FileImporter
}

type jsonValue struct {
	val interface{}
	// these objects are exclusively owned by this jsonValue
	owned []*C.struct_JsonnetJsonValue
}

var handles = handlesTable{}
var versionString *C.char

//export jsonnet_version
func jsonnet_version() *C.char {
	if versionString == nil {
		version := jsonnet.Version() + " (go-jsonnet)"
		versionString = C.CString(version)
	}
	return versionString
}

//export jsonnet_make
func jsonnet_make() *C.struct_JsonnetVm {
	newVM := &vm{jsonnet.MakeVM(), &jsonnet.FileImporter{}}
	newVM.Importer(newVM.importer)

	id, err := handles.make(newVM)

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	return C.jsonnet_internal_make_vm_with_id(C.uint32_t(id))
}

//export jsonnet_destroy
func jsonnet_destroy(vmRef *C.struct_JsonnetVm) {
	if err := handles.free(uint32(vmRef.id)); err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	C.jsonnet_internal_free_vm(vmRef)
}

func getVM(vmRef *C.struct_JsonnetVm) *vm {
	ref, err := handles.get(uint32(vmRef.id))

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	v, ok := ref.(*vm)

	if !ok {
		fmt.Fprintf(os.Stderr, "provided handle has a different type")
		os.Exit(1)
	}

	return v
}

func evaluateSnippet(vmRef *C.struct_JsonnetVm, filename string, code string, e *C.int) *C.char {
	vm := getVM(vmRef)
	out, err := vm.EvaluateSnippet(filename, code)
	var result *C.char
	if err != nil {
		*e = 1
		result = C.CString(err.Error())
	} else {
		*e = 0
		result = C.CString(out)
	}
	return result
}

//export jsonnet_evaluate_snippet
func jsonnet_evaluate_snippet(vmRef *C.struct_JsonnetVm, filename *C.char, code *C.char, e *C.int) *C.char {
	f := C.GoString(filename)
	s := C.GoString(code)
	return evaluateSnippet(vmRef, f, s, e)
}

//export jsonnet_evaluate_file
func jsonnet_evaluate_file(vmRef *C.struct_JsonnetVm, filename *C.char, e *C.int) *C.char {
	f := C.GoString(filename)
	data, err := ioutil.ReadFile(f)
	if err != nil {
		*e = 1
		// TODO(sbarzowski) make sure that it's ok allocation-wise
		return C.CString(fmt.Sprintf("Failed to read input file: %s: %s", f, err.Error()))
	}
	return evaluateSnippet(vmRef, f, string(data), e)
}

//export jsonnet_max_stack
func jsonnet_max_stack(vmRef *C.struct_JsonnetVm, v C.uint) {
	vm := getVM(vmRef)
	vm.MaxStack = int(v) // potentially dangerous conversion
}

//export jsonnet_string_output
func jsonnet_string_output(vmRef *C.struct_JsonnetVm, v C.int) {
	vm := getVM(vmRef)
	vm.StringOutput = v != 0
}

//export jsonnet_max_trace
func jsonnet_max_trace(vmRef *C.struct_JsonnetVm, v C.uint) {
	vm := getVM(vmRef)
	vm.ErrorFormatter.SetMaxStackTraceSize(int(v)) // potentially dangerous conversion
}

//export jsonnet_jpath_add
func jsonnet_jpath_add(vmRef *C.struct_JsonnetVm, path *C.char) {
	vm := getVM(vmRef)
	vm.importer.JPaths = append(vm.importer.JPaths, C.GoString(path))
}

//export jsonnet_ext_var
func jsonnet_ext_var(vmRef *C.struct_JsonnetVm, key, value *C.char) {
	vm := getVM(vmRef)
	vm.ExtVar(C.GoString(key), C.GoString(value))
}

//export jsonnet_ext_code
func jsonnet_ext_code(vmRef *C.struct_JsonnetVm, key, value *C.char) {
	vm := getVM(vmRef)
	vm.ExtCode(C.GoString(key), C.GoString(value))
}

//export jsonnet_tla_var
func jsonnet_tla_var(vmRef *C.struct_JsonnetVm, key, value *C.char) {
	vm := getVM(vmRef)
	vm.TLAVar(C.GoString(key), C.GoString(value))
}

//export jsonnet_tla_code
func jsonnet_tla_code(vmRef *C.struct_JsonnetVm, key, value *C.char) {
	vm := getVM(vmRef)
	vm.TLACode(C.GoString(key), C.GoString(value))
}

//export jsonnet_native_callback
func jsonnet_native_callback(vmRef *C.struct_JsonnetVm, name *C.char, cb *C.JsonnetNativeCallback, ctx unsafe.Pointer, params **C.char) {
	vm := getVM(vmRef)
	p := unsafe.Pointer(params)
	sz := unsafe.Sizeof(*params)

	var paramNames ast.Identifiers

	for i := 0; ; i++ {
		param := (**C.char)(unsafe.Pointer(uintptr(p) + uintptr(i)*sz))

		if *param == nil {
			break
		}

		paramNames = append(paramNames, ast.Identifier(C.GoString(*param)))
	}

	f := &jsonnet.NativeFunction{
		Name:   C.GoString(name),
		Params: paramNames,
		Func: func(x []interface{}) (interface{}, error) {
			var (
				arr     []*C.struct_JsonnetJsonValue
				argv    **C.struct_JsonnetJsonValue
				success = C.int(0)
			)

			if len(x) > 0 {
				arr = make([]*C.struct_JsonnetJsonValue, 0, len(x))

				for _, json := range x {
					arr = append(arr, createJSONValue(vmRef, json))
				}

				argv = &(arr[0])
			}

			result := C.jsonnet_internal_execute_native(cb, ctx, argv, &success)
			v := getJSONValue(result)

			for _, val := range arr {
				jsonnet_json_destroy(vmRef, val)
			}

			jsonnet_json_destroy(vmRef, result)

			if success != 1 {
				return nil, fmt.Errorf("failed to execute native callback, code: %d", success)
			}

			return v.val, nil
		},
	}

	vm.NativeFunction(f)
}

//export jsonnet_json_extract_string
func jsonnet_json_extract_string(vmRef *C.struct_JsonnetVm, json *C.struct_JsonnetJsonValue) *C.char {
	v := getJSONValue(json)
	str, ok := v.val.(string)

	if !ok {
		return nil
	}

	return C.CString(str)
}

//export jsonnet_json_extract_number
func jsonnet_json_extract_number(vmRef *C.struct_JsonnetVm, json *C.struct_JsonnetJsonValue, out *C.double) C.int {
	v := getJSONValue(json)

	switch f := v.val.(type) {
	case float64:
		*out = C.double(f)
		return 1
	case int:
		*out = C.double(f)
		return 1
	}

	return 0
}

//export jsonnet_json_extract_bool
func jsonnet_json_extract_bool(vmRef *C.struct_JsonnetVm, json *C.struct_JsonnetJsonValue) C.int {
	v := getJSONValue(json)
	b, ok := v.val.(bool)

	if !ok {
		return 2
	}

	if b {
		return 1
	}

	return 0
}

//export jsonnet_json_extract_null
func jsonnet_json_extract_null(vmRef *C.struct_JsonnetVm, json *C.struct_JsonnetJsonValue) C.int {
	v := getJSONValue(json)

	if v.val == nil {
		return 1
	}

	return 0
}

//export jsonnet_json_make_string
func jsonnet_json_make_string(vmRef *C.struct_JsonnetVm, v *C.char) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, C.GoString(v))
}

//export jsonnet_json_make_number
func jsonnet_json_make_number(vmRef *C.struct_JsonnetVm, v C.double) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, float64(v))
}

//export jsonnet_json_make_bool
func jsonnet_json_make_bool(vmRef *C.struct_JsonnetVm, v C.int) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, v != 0)
}

//export jsonnet_json_make_null
func jsonnet_json_make_null(vmRef *C.struct_JsonnetVm) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, nil)
}

//export jsonnet_json_make_array
func jsonnet_json_make_array(vmRef *C.struct_JsonnetVm) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, []interface{}{})
}

//export jsonnet_json_array_append
func jsonnet_json_array_append(vmRef *C.struct_JsonnetVm, arr *C.struct_JsonnetJsonValue, v *C.struct_JsonnetJsonValue) {
	json := getJSONValue(arr)
	slice, ok := json.val.([]interface{})

	if !ok {
		fmt.Fprintf(os.Stderr, "array should be provided")
		os.Exit(1)
	}

	json.val = append(slice, getJSONValue(v).val)
	json.owned = append(json.owned, v)
}

//export jsonnet_json_make_object
func jsonnet_json_make_object(vmRef *C.struct_JsonnetVm) *C.struct_JsonnetJsonValue {
	return createJSONValue(vmRef, make(map[string]interface{}))
}

//export jsonnet_json_object_append
func jsonnet_json_object_append(
	vmRef *C.struct_JsonnetVm,
	obj *C.struct_JsonnetJsonValue,
	f *C.char,
	v *C.struct_JsonnetJsonValue,
) {
	d := getJSONValue(obj)
	table, ok := d.val.(map[string]interface{})

	if !ok {
		fmt.Fprintf(os.Stderr, "object should be provided")
		os.Exit(1)
	}

	table[C.GoString(f)] = getJSONValue(v).val
	d.owned = append(d.owned, v)
}

//export jsonnet_json_destroy
func jsonnet_json_destroy(vmRef *C.struct_JsonnetVm, v *C.struct_JsonnetJsonValue) {
	for _, child := range getJSONValue(v).owned {
		jsonnet_json_destroy(vmRef, child)
	}

	if err := handles.free(uint32(v.id)); err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	C.jsonnet_internal_free_json(v)
}

func createJSONValue(vmRef *C.struct_JsonnetVm, val interface{}) *C.struct_JsonnetJsonValue {
	id, err := handles.make(&jsonValue{val: val})

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	return C.jsonnet_internal_make_json_with_id(C.uint32_t(id))
}

func getJSONValue(jsonRef *C.struct_JsonnetJsonValue) *jsonValue {
	ref, err := handles.get(uint32(jsonRef.id))

	if err != nil {
		fmt.Fprintf(os.Stderr, err.Error())
		os.Exit(1)
	}

	v, ok := ref.(*jsonValue)

	if !ok {
		fmt.Fprintf(os.Stderr, "provided handle has a different type")
		os.Exit(1)
	}

	return v
}

func main() {
}
<|endoftext|>