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

package apiserver

import (
	"crypto/tls"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/loggo/loggo"
	"launchpad.net/tomb"

	"launchpad.net/juju-core/rpc"
	"launchpad.net/juju-core/rpc/jsoncodec"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/apiserver/common"
)

var logger = loggo.GetLogger("juju.state.apiserver")

// Server holds the server side of the API.
type Server struct {
	tomb    tomb.Tomb
	wg      sync.WaitGroup
	state   *state.State
	addr    net.Addr
	dataDir string
}

// Serve serves the given state by accepting requests on the given
// listener, using the given certificate and key (in PEM format) for
// authentication.
func NewServer(s *state.State, addr string, cert, key []byte, datadir string) (*Server, error) {
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	logger.Infof("listening on %q", lis.Addr())
	tlsCert, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return nil, err
	}
	srv := &Server{
		state:   s,
		addr:    lis.Addr(),
		dataDir: datadir,
	}
	// TODO(rog) check that *srvRoot is a valid type for using
	// as an RPC server.
	lis = tls.NewListener(lis, &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
	})
	go srv.run(lis)
	return srv, nil
}

// Dead returns a channel that signals when the server has exited.
func (srv *Server) Dead() <-chan struct{} {
	return srv.tomb.Dead()
}

// Stop stops the server and returns when all running requests
// have completed.
func (srv *Server) Stop() error {
	srv.tomb.Kill(nil)
	return srv.tomb.Wait()
}

// Kill implements worker.Worker.Kill.
func (srv *Server) Kill() {
	srv.tomb.Kill(nil)
}

// Wait implements worker.Worker.Wait.
func (srv *Server) Wait() error {
	return srv.tomb.Wait()
}

type requestNotifier struct {
	id    int64
	start time.Time

	mu   sync.Mutex
	tag_ string
}

var globalCounter int64

func newRequestNotifier() *requestNotifier {
	return &requestNotifier{
		id:    atomic.AddInt64(&globalCounter, 1),
		tag_:  "<unknown>",
		start: time.Now(),
	}
}

func (n *requestNotifier) login(tag string) {
	n.mu.Lock()
	n.tag_ = tag
	n.mu.Unlock()
}

func (n *requestNotifier) tag() (tag string) {
	n.mu.Lock()
	tag = n.tag_
	n.mu.Unlock()
	return
}

func (n *requestNotifier) ServerRequest(hdr *rpc.Header, body interface{}) {
	if hdr.Request.Type == "Pinger" && hdr.Request.Action == "Ping" {
		return
	}
	// TODO(rog) 2013-10-11 remove secrets from some requests.
	logger.Debugf("<- [%X] %s %s", n.id, n.tag(), jsoncodec.DumpRequest(hdr, body))
}

func (n *requestNotifier) ServerReply(req rpc.Request, hdr *rpc.Header, body interface{}, timeSpent time.Duration) {
	if req.Type == "Pinger" && req.Action == "Ping" {
		return
	}
	logger.Debugf("-> [%X] %s %s %s %s[%q].%s", n.id, n.tag(), timeSpent, jsoncodec.DumpRequest(hdr, body), req.Type, req.Id, req.Action)
}

func (n *requestNotifier) join(req *http.Request) {
	logger.Infof("[%X] API connection from %s", n.id, req.RemoteAddr)
}

func (n *requestNotifier) leave() {
	logger.Infof("[%X] API connection terminated after %v", n.id, time.Since(n.start))
}

func (n requestNotifier) ClientRequest(hdr *rpc.Header, body interface{}) {
}

func (n requestNotifier) ClientReply(req rpc.Request, hdr *rpc.Header, body interface{}) {
}

func (srv *Server) run(lis net.Listener) {
	defer srv.tomb.Done()
	defer srv.wg.Wait() // wait for any outstanding requests to complete.
	srv.wg.Add(1)
	go func() {
		<-srv.tomb.Dying()
		lis.Close()
		srv.wg.Done()
	}()
	mux := http.NewServeMux()
	mux.HandleFunc("/", srv.apiHandler)
	mux.Handle("/charms", &charmsHandler{state: srv.state})
	// The error from http.Serve is not interesting.
	http.Serve(lis, mux)
}

func (srv *Server) apiHandler(w http.ResponseWriter, req *http.Request) {
	reqNotifier := newRequestNotifier()
	reqNotifier.join(req)
	defer reqNotifier.leave()
	wsServer := websocket.Server{
		Handler: func(conn *websocket.Conn) {
			srv.wg.Add(1)
			defer srv.wg.Done()
			// If we've got to this stage and the tomb is still
			// alive, we know that any tomb.Kill must occur after we
			// have called wg.Add, so we avoid the possibility of a
			// handler goroutine running after Stop has returned.
			if srv.tomb.Err() != tomb.ErrStillAlive {
				return
			}
			if err := srv.serveConn(conn, reqNotifier); err != nil {
				logger.Errorf("error serving RPCs: %v", err)
			}
		},
	}
	wsServer.ServeHTTP(w, req)
}

// Addr returns the address that the server is listening on.
func (srv *Server) Addr() string {
	return srv.addr.String()
}

func (srv *Server) serveConn(wsConn *websocket.Conn, reqNotifier *requestNotifier) error {
	codec := jsoncodec.NewWebsocket(wsConn)
	if loggo.GetLogger("juju.rpc.jsoncodec").EffectiveLogLevel() <= loggo.TRACE {
		codec.SetLogging(true)
	}
	var notifier rpc.RequestNotifier
	if logger.EffectiveLogLevel() <= loggo.DEBUG {
		// Incur request monitoring overhead only if we
		// know we'll need it.
		notifier = reqNotifier
	}
	conn := rpc.NewConn(codec, notifier)
	conn.Serve(newStateServer(srv, conn, reqNotifier), serverError)
	conn.Start()
	select {
	case <-conn.Dead():
	case <-srv.tomb.Dying():
	}
	return conn.Close()
}

func serverError(err error) error {
	if err := common.ServerError(err); err != nil {
		return err
	}
	return nil
}

var logRequests = true
<commit_msg>Add the tag for disconnect<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package apiserver

import (
	"crypto/tls"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"

	"code.google.com/p/go.net/websocket"
	"github.com/loggo/loggo"
	"launchpad.net/tomb"

	"launchpad.net/juju-core/rpc"
	"launchpad.net/juju-core/rpc/jsoncodec"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/apiserver/common"
)

var logger = loggo.GetLogger("juju.state.apiserver")

// Server holds the server side of the API.
type Server struct {
	tomb    tomb.Tomb
	wg      sync.WaitGroup
	state   *state.State
	addr    net.Addr
	dataDir string
}

// Serve serves the given state by accepting requests on the given
// listener, using the given certificate and key (in PEM format) for
// authentication.
func NewServer(s *state.State, addr string, cert, key []byte, datadir string) (*Server, error) {
	lis, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	logger.Infof("listening on %q", lis.Addr())
	tlsCert, err := tls.X509KeyPair(cert, key)
	if err != nil {
		return nil, err
	}
	srv := &Server{
		state:   s,
		addr:    lis.Addr(),
		dataDir: datadir,
	}
	// TODO(rog) check that *srvRoot is a valid type for using
	// as an RPC server.
	lis = tls.NewListener(lis, &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
	})
	go srv.run(lis)
	return srv, nil
}

// Dead returns a channel that signals when the server has exited.
func (srv *Server) Dead() <-chan struct{} {
	return srv.tomb.Dead()
}

// Stop stops the server and returns when all running requests
// have completed.
func (srv *Server) Stop() error {
	srv.tomb.Kill(nil)
	return srv.tomb.Wait()
}

// Kill implements worker.Worker.Kill.
func (srv *Server) Kill() {
	srv.tomb.Kill(nil)
}

// Wait implements worker.Worker.Wait.
func (srv *Server) Wait() error {
	return srv.tomb.Wait()
}

type requestNotifier struct {
	id    int64
	start time.Time

	mu   sync.Mutex
	tag_ string
}

var globalCounter int64

func newRequestNotifier() *requestNotifier {
	return &requestNotifier{
		id:    atomic.AddInt64(&globalCounter, 1),
		tag_:  "<unknown>",
		start: time.Now(),
	}
}

func (n *requestNotifier) login(tag string) {
	n.mu.Lock()
	n.tag_ = tag
	n.mu.Unlock()
}

func (n *requestNotifier) tag() (tag string) {
	n.mu.Lock()
	tag = n.tag_
	n.mu.Unlock()
	return
}

func (n *requestNotifier) ServerRequest(hdr *rpc.Header, body interface{}) {
	if hdr.Request.Type == "Pinger" && hdr.Request.Action == "Ping" {
		return
	}
	// TODO(rog) 2013-10-11 remove secrets from some requests.
	logger.Debugf("<- [%X] %s %s", n.id, n.tag(), jsoncodec.DumpRequest(hdr, body))
}

func (n *requestNotifier) ServerReply(req rpc.Request, hdr *rpc.Header, body interface{}, timeSpent time.Duration) {
	if req.Type == "Pinger" && req.Action == "Ping" {
		return
	}
	logger.Debugf("-> [%X] %s %s %s %s[%q].%s", n.id, n.tag(), timeSpent, jsoncodec.DumpRequest(hdr, body), req.Type, req.Id, req.Action)
}

func (n *requestNotifier) join(req *http.Request) {
	logger.Infof("[%X] API connection from %s", n.id, req.RemoteAddr)
}

func (n *requestNotifier) leave() {
	logger.Infof("[%X] %s API connection terminated after %v", n.id, n.tag(), time.Since(n.start))
}

func (n requestNotifier) ClientRequest(hdr *rpc.Header, body interface{}) {
}

func (n requestNotifier) ClientReply(req rpc.Request, hdr *rpc.Header, body interface{}) {
}

func (srv *Server) run(lis net.Listener) {
	defer srv.tomb.Done()
	defer srv.wg.Wait() // wait for any outstanding requests to complete.
	srv.wg.Add(1)
	go func() {
		<-srv.tomb.Dying()
		lis.Close()
		srv.wg.Done()
	}()
	mux := http.NewServeMux()
	mux.HandleFunc("/", srv.apiHandler)
	mux.Handle("/charms", &charmsHandler{state: srv.state})
	// The error from http.Serve is not interesting.
	http.Serve(lis, mux)
}

func (srv *Server) apiHandler(w http.ResponseWriter, req *http.Request) {
	reqNotifier := newRequestNotifier()
	reqNotifier.join(req)
	defer reqNotifier.leave()
	wsServer := websocket.Server{
		Handler: func(conn *websocket.Conn) {
			srv.wg.Add(1)
			defer srv.wg.Done()
			// If we've got to this stage and the tomb is still
			// alive, we know that any tomb.Kill must occur after we
			// have called wg.Add, so we avoid the possibility of a
			// handler goroutine running after Stop has returned.
			if srv.tomb.Err() != tomb.ErrStillAlive {
				return
			}
			if err := srv.serveConn(conn, reqNotifier); err != nil {
				logger.Errorf("error serving RPCs: %v", err)
			}
		},
	}
	wsServer.ServeHTTP(w, req)
}

// Addr returns the address that the server is listening on.
func (srv *Server) Addr() string {
	return srv.addr.String()
}

func (srv *Server) serveConn(wsConn *websocket.Conn, reqNotifier *requestNotifier) error {
	codec := jsoncodec.NewWebsocket(wsConn)
	if loggo.GetLogger("juju.rpc.jsoncodec").EffectiveLogLevel() <= loggo.TRACE {
		codec.SetLogging(true)
	}
	var notifier rpc.RequestNotifier
	if logger.EffectiveLogLevel() <= loggo.DEBUG {
		// Incur request monitoring overhead only if we
		// know we'll need it.
		notifier = reqNotifier
	}
	conn := rpc.NewConn(codec, notifier)
	conn.Serve(newStateServer(srv, conn, reqNotifier), serverError)
	conn.Start()
	select {
	case <-conn.Dead():
	case <-srv.tomb.Dying():
	}
	return conn.Close()
}

func serverError(err error) error {
	if err := common.ServerError(err); err != nil {
		return err
	}
	return nil
}

var logRequests = true
<|endoftext|>


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

package files

import (
	"path/filepath"

	"github.com/juju/errors"
)

// TODO(ericsnow) Pull these from authoritative sources (see
// github.com/juju/juju/juju/paths, etc.):
const (
	dataDir        = "/var/lib/juju"
	startupDir     = "/etc/init"
	loggingConfDir = "/etc/rsyslog.d"
	logsDir        = "/var/log/juju"
	sshDir         = "/home/ubuntu/.ssh"

	machinesConfs = "jujud-machine-*.conf"
	agentsDir     = "agents"
	agentsConfs   = "machine-*"
	loggingConfs  = "*juju.conf"
	toolsDir      = "tools"

	sshIdentFile   = "system-identity"
	nonceFile      = "nonce.txt"
	allMachinesLog = "all-machines.log"
	machine0Log    = "machine-0.log"
	authKeysFile   = "authorized_keys"

	dbStartupConf = "juju-db.conf"
	dbPEM         = "server.pem"
	dbSecret      = "shared-secret"
)

// Paths holds the paths that backups needs.
type Paths struct {
	DataDir string
	LogsDir string
}

// GetFilesToBackUp returns the paths that should be included in the
// backup archive.
func GetFilesToBackUp(rootDir string, paths Paths) ([]string, error) {
	var glob string

	glob = filepath.Join(rootDir, startupDir, machinesConfs)
	initMachineConfs, err := filepath.Glob(glob)
	if err != nil {
		return nil, errors.Annotate(err, "failed to fetch machine init files")
	}

	glob = filepath.Join(rootDir, paths.DataDir, agentsDir, agentsConfs)
	agentConfs, err := filepath.Glob(glob)
	if err != nil {
		return nil, errors.Annotate(err, "failed to fetch agent config files")
	}

	glob = filepath.Join(rootDir, loggingConfDir, loggingConfs)
	jujuLogConfs, err := filepath.Glob(glob)
	if err != nil {
		return nil, errors.Annotate(err, "failed to fetch juju log conf files")
	}

	backupFiles := []string{
		filepath.Join(rootDir, paths.DataDir, toolsDir),

		filepath.Join(rootDir, paths.DataDir, sshIdentFile),
		filepath.Join(rootDir, paths.DataDir, nonceFile),
		filepath.Join(rootDir, paths.LogsDir, allMachinesLog),
		filepath.Join(rootDir, paths.LogsDir, machine0Log),
		filepath.Join(rootDir, sshDir, authKeysFile),

		filepath.Join(rootDir, paths.DataDir, dbPEM),
		filepath.Join(rootDir, paths.DataDir, dbSecret),
	}
	backupFiles = append(backupFiles, initMachineConfs...)
	backupFiles = append(backupFiles, agentConfs...)
	backupFiles = append(backupFiles, jujuLogConfs...)

	return backupFiles, nil
}
<commit_msg>Accommodate a missing .ssh directory.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package files

import (
	"os"
	"path/filepath"

	"github.com/juju/errors"
)

// TODO(ericsnow) Pull these from authoritative sources (see
// github.com/juju/juju/juju/paths, etc.):
const (
	dataDir        = "/var/lib/juju"
	startupDir     = "/etc/init"
	loggingConfDir = "/etc/rsyslog.d"
	logsDir        = "/var/log/juju"
	sshDir         = "/home/ubuntu/.ssh"

	machinesConfs = "jujud-machine-*.conf"
	agentsDir     = "agents"
	agentsConfs   = "machine-*"
	loggingConfs  = "*juju.conf"
	toolsDir      = "tools"

	sshIdentFile   = "system-identity"
	nonceFile      = "nonce.txt"
	allMachinesLog = "all-machines.log"
	machine0Log    = "machine-0.log"
	authKeysFile   = "authorized_keys"

	dbStartupConf = "juju-db.conf"
	dbPEM         = "server.pem"
	dbSecret      = "shared-secret"
)

// Paths holds the paths that backups needs.
type Paths struct {
	DataDir string
	LogsDir string
}

// GetFilesToBackUp returns the paths that should be included in the
// backup archive.
func GetFilesToBackUp(rootDir string, paths Paths) ([]string, error) {
	var glob string

	glob = filepath.Join(rootDir, startupDir, machinesConfs)
	initMachineConfs, err := filepath.Glob(glob)
	if err != nil {
		return nil, errors.Annotate(err, "failed to fetch machine init files")
	}

	glob = filepath.Join(rootDir, paths.DataDir, agentsDir, agentsConfs)
	agentConfs, err := filepath.Glob(glob)
	if err != nil {
		return nil, errors.Annotate(err, "failed to fetch agent config files")
	}

	glob = filepath.Join(rootDir, loggingConfDir, loggingConfs)
	jujuLogConfs, err := filepath.Glob(glob)
	if err != nil {
		return nil, errors.Annotate(err, "failed to fetch juju log conf files")
	}

	backupFiles := []string{
		filepath.Join(rootDir, paths.DataDir, toolsDir),

		filepath.Join(rootDir, paths.DataDir, sshIdentFile),
		filepath.Join(rootDir, paths.DataDir, nonceFile),
		filepath.Join(rootDir, paths.LogsDir, allMachinesLog),
		filepath.Join(rootDir, paths.LogsDir, machine0Log),

		filepath.Join(rootDir, paths.DataDir, dbPEM),
		filepath.Join(rootDir, paths.DataDir, dbSecret),
	}
	backupFiles = append(backupFiles, initMachineConfs...)
	backupFiles = append(backupFiles, agentConfs...)
	backupFiles = append(backupFiles, jujuLogConfs...)

	// Handle user SSH files (might not exist).
	SSHDir := filepath.Join(rootDir, sshDir)
	if _, err := os.Stat(SSHDir); err != nil {
		if !os.IsNotExist(err) {
			return nil, errors.Trace(err)
		}
	} else {
		backupFiles = append(backupFiles, filepath.Join(SSHDir, authKeysFile))
	}

	return backupFiles, nil
}
<|endoftext|>


<commit_before>// Package magick is a hacked up port of the minimal functionality we need
// to satisfy the IIIFImageDecoder interface.  Code is based in part on
// github.com/quirkey/magick
package magick

/*
#cgo pkg-config: MagickCore
#include <magick/MagickCore.h>
*/
import "C"
import (
	"fmt"
	"os"
	"unsafe"
)

func init() {
	path, _ := os.Getwd()
	cPath := C.CString(path)
	defer C.free(unsafe.Pointer(cPath))
	C.MagickCoreGenesis(cPath, C.MagickFalse)
}

func makeError(exception *C.ExceptionInfo) error {
	return fmt.Errorf("%s: %s - %s", exception.severity, exception.reason, exception.description)
}
<commit_msg>magick: use %v for formatting C data<commit_after>// Package magick is a hacked up port of the minimal functionality we need
// to satisfy the IIIFImageDecoder interface.  Code is based in part on
// github.com/quirkey/magick
package magick

/*
#cgo pkg-config: MagickCore
#include <magick/MagickCore.h>
*/
import "C"
import (
	"fmt"
	"os"
	"unsafe"
)

func init() {
	path, _ := os.Getwd()
	cPath := C.CString(path)
	defer C.free(unsafe.Pointer(cPath))
	C.MagickCoreGenesis(cPath, C.MagickFalse)
}

func makeError(exception *C.ExceptionInfo) error {
	return fmt.Errorf("%v: %v - %v", exception.severity, exception.reason, exception.description)
}
<|endoftext|>


<commit_before>package utils

import (
	"bytes"
	"text/template"
	"os/exec"
	"syscall"
	"fmt"
	"github.com/pkg/errors"
	"github.com/Sirupsen/logrus"
	"io/ioutil"
	"os"
)

type Bash struct {
	Command string
	PipeFail bool
	Arguments map[string]string
	NoLog bool

	retCode int
	stdout string
	stderr string
	err error
}

func (b *Bash) build() error {
	Assert(b.Command != "", "Command cannot be emptry string")

	if (b.Arguments != nil) {
		tmpl, err := template.New("script").Parse(b.Command)
		if err != nil {
			return err
		}

		var buf bytes.Buffer
		err = tmpl.Execute(&buf, b.Arguments)
		if err != nil {
			return err
		}

		b.Command = buf.String()
	}

	if b.PipeFail {
		b.Command = fmt.Sprintf("set -o pipefail; %s", b.Command)
	}

	return nil
}

func (b *Bash) Run() error {
	ret, so, se, err := b.RunWithReturn()
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to execute the command[%s] because of an internal errro",  b.Command))
	}

	if ret != 0 {
		return errors.New(fmt.Sprintf("failed to exectue the command[%s]\nreturn code:%d\nstdout:%s\nstderr:%s\n",
			b.Command, ret, so, se))
	}

	return nil
}

func (b *Bash) RunWithReturn() (retCode int, stdout, stderr string, err error) {
	if err = b.build(); err != nil {
		b.err = err
		return -1, "", "", err
	}

	if !b.NoLog {
		logrus.Debugf("shell start: %s", b.Command)
	}

	var so, se bytes.Buffer
	var cmd *exec.Cmd

	if len(b.Command) > 1024* 4 {
		content := []byte(b.Command)
		tmpfile, err := ioutil.TempFile("/home/vyos", "zvrcommand"); PanicOnError(err)
		ioutil.WriteFile(tmpfile.Name(), content, 0777)
		tmpfile.Close()
		cmd = exec.Command("bash", "-c", tmpfile.Name())
		//path := "/home/vyos/zvrcommand"
		//ioutil.WriteFile(path, content, 0777)
		//cmd = exec.Command("bash", "-c", path)
		defer os.Remove(tmpfile.Name())
	} else {
		cmd = exec.Command("bash", "-c", b.Command)
	}

	cmd.Stdout = &so
	cmd.Stderr = &se

	var waitStatus syscall.WaitStatus
	if err := cmd.Run(); err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			waitStatus = exitError.Sys().(syscall.WaitStatus)
			retCode = waitStatus.ExitStatus()
		} else {
			panic(errors.Errorf("unable to get return code, %s", err))
		}
	} else {
		waitStatus = cmd.ProcessState.Sys().(syscall.WaitStatus)
		retCode = waitStatus.ExitStatus()
	}

	stdout = string(so.Bytes())
	stderr = string(se.Bytes())

	b.retCode = retCode
	b.stdout = stdout
	b.stderr = stderr

	if !b.NoLog {
		logrus.WithFields(logrus.Fields{
			"return code": fmt.Sprintf("%v", retCode),
			"stdout": stdout,
			"stderr": stderr,
		}).Debugf("shell done: %s", b.Command)
	}

	return
}

func (bash *Bash) PanicIfError() {
	if bash.err != nil {
		panic(errors.New(fmt.Sprintf("shell failure[command: %v], internal error: %v",
			bash.Command, bash.err)))
	}

	if bash.retCode != 0 {
		panic(errors.New(fmt.Sprintf("shell failure[command: %v, return code: %v, stdout: %v, stderr: %v",
			bash.Command, bash.retCode, bash.stdout, bash.stderr)))
	}
}

func NewBash() *Bash {
	return &Bash{}
}

<commit_msg>give exec permission to script<commit_after>package utils

import (
	"bytes"
	"text/template"
	"os/exec"
	"syscall"
	"fmt"
	"github.com/pkg/errors"
	"github.com/Sirupsen/logrus"
	"io/ioutil"
	//"os"
	"os"
)

type Bash struct {
	Command string
	PipeFail bool
	Arguments map[string]string
	NoLog bool

	retCode int
	stdout string
	stderr string
	err error
}

func (b *Bash) build() error {
	Assert(b.Command != "", "Command cannot be emptry string")

	if (b.Arguments != nil) {
		tmpl, err := template.New("script").Parse(b.Command)
		if err != nil {
			return err
		}

		var buf bytes.Buffer
		err = tmpl.Execute(&buf, b.Arguments)
		if err != nil {
			return err
		}

		b.Command = buf.String()
	}

	if b.PipeFail {
		b.Command = fmt.Sprintf("set -o pipefail; %s", b.Command)
	}

	return nil
}

func (b *Bash) Run() error {
	ret, so, se, err := b.RunWithReturn()
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to execute the command[%s] because of an internal errro",  b.Command))
	}

	if ret != 0 {
		return errors.New(fmt.Sprintf("failed to exectue the command[%s]\nreturn code:%d\nstdout:%s\nstderr:%s\n",
			b.Command, ret, so, se))
	}

	return nil
}

func (b *Bash) RunWithReturn() (retCode int, stdout, stderr string, err error) {
	if err = b.build(); err != nil {
		b.err = err
		return -1, "", "", err
	}

	if !b.NoLog {
		logrus.Debugf("shell start: %s", b.Command)
	}

	var so, se bytes.Buffer
	var cmd *exec.Cmd

	if len(b.Command) > 1024* 4 {
		content := []byte(b.Command)
		tmpfile, err := ioutil.TempFile("/home/vyos", "zvrcommand"); PanicOnError(err)
		err = tmpfile.Chmod(0777); PanicOnError(err)
		_, err = tmpfile.Write(content); PanicOnError(err)
		err = tmpfile.Close(); PanicOnError(err)
		cmd = exec.Command("bash", "-c", tmpfile.Name())
		//ioutil.WriteFile(tmpfile.Name(), content, 0777)
		//path := "/home/vyos/zvrcommand"
		//ioutil.WriteFile(path, content, 0777)
		//cmd = exec.Command("bash", "-c", path)
		defer os.Remove(tmpfile.Name())
	} else {
		cmd = exec.Command("bash", "-c", b.Command)
	}

	cmd.Stdout = &so
	cmd.Stderr = &se

	var waitStatus syscall.WaitStatus
	if err := cmd.Run(); err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			waitStatus = exitError.Sys().(syscall.WaitStatus)
			retCode = waitStatus.ExitStatus()
		} else {
			panic(errors.Errorf("unable to get return code, %s", err))
		}
	} else {
		waitStatus = cmd.ProcessState.Sys().(syscall.WaitStatus)
		retCode = waitStatus.ExitStatus()
	}

	stdout = string(so.Bytes())
	stderr = string(se.Bytes())

	b.retCode = retCode
	b.stdout = stdout
	b.stderr = stderr

	if !b.NoLog {
		logrus.WithFields(logrus.Fields{
			"return code": fmt.Sprintf("%v", retCode),
			"stdout": stdout,
			"stderr": stderr,
		}).Debugf("shell done: %s", b.Command)
	}

	return
}

func (bash *Bash) PanicIfError() {
	if bash.err != nil {
		panic(errors.New(fmt.Sprintf("shell failure[command: %v], internal error: %v",
			bash.Command, bash.err)))
	}

	if bash.retCode != 0 {
		panic(errors.New(fmt.Sprintf("shell failure[command: %v, return code: %v, stdout: %v, stderr: %v",
			bash.Command, bash.retCode, bash.stdout, bash.stderr)))
	}
}

func NewBash() *Bash {
	return &Bash{}
}

<|endoftext|>


<commit_before>package library

import (
	"encoding/json"
	"log"

	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/sqs"

	"strconv"
	"sync"

	"github.com/nytlabs/streamtools/st/blocks" // blocks
	"github.com/nytlabs/streamtools/st/util"
)

// lots of this code stolen brazenly from JP https://github.com/jprobinson
func (b *FromSQS) listener() {
	accessKey, ok := b.auth["AccessKey"].(string)
	if !ok {
		b.Error("could not assert AccessKey to a string")
		return
	}
	accessSecret, ok := b.auth["AccessSecret"].(string)
	if !ok {
		b.Error("could not assert AccessSecret to a string")
		return
	}
	queueName, ok := b.auth["QueueName"].(string)
	if !ok {
		b.Error("could not assert queue name to a string")
		return
	}
	auth := aws.Auth{AccessKey: accessKey, SecretKey: accessSecret}
	sqsClient := sqs.New(auth, aws.USEast)
	queue, err := sqsClient.GetQueue(queueName)

	maxNstring, ok := b.auth["MaxNumberOfMessages"].(string)
	if !ok {
		b.Error("could not assert MaxNumberOfMessages to a string")
		return
	}
	maxN, err := strconv.Atoi(maxNstring)
	if err != nil {
		b.Error(err)
		return
	}

	var resp *sqs.ReceiveMessageResponse
	for {
		select {
		case <-b.stop:
			log.Println("Exiting SQS read loop")
			return
		default:
			resp, err = queue.ReceiveMessage(maxN)
			if err != nil {
				b.Error(err)
			}
			if len(resp.Messages) == 0 {
				break
			}
			for _, m := range resp.Messages {
				select {
				case b.fromListener <- []byte(m.Body):
				default:
					log.Println("discarding messages")
					log.Println(len(b.fromListener))
					continue
				}

				if _, err = queue.DeleteMessage(&m); err != nil {
					b.Error(err)
				}
			}
		}
	}
}

// specify those channels we're going to use to communicate with streamtools
type FromSQS struct {
	blocks.Block
	queryrule chan blocks.MsgChan
	inrule    blocks.MsgChan
	out       blocks.MsgChan
	quit      blocks.MsgChan

	lock         sync.Mutex
	listening    bool
	fromListener chan []byte
	auth         map[string]interface{}
	stop         chan bool
}

// we need to build a simple factory so that streamtools can make new blocks of this kind
func NewFromSQS() blocks.BlockInterface {
	return &FromSQS{}
}

// Setup is called once before running the block. We build up the channels and specify what kind of block this is.
func (b *FromSQS) Setup() {
	b.Kind = "Queue I/O"
	b.Desc = "reads from Amazon's SQS, emitting each line of JSON as a separate message"
	b.inrule = b.InRoute("rule")
	b.queryrule = b.QueryRoute("rule")
	b.quit = b.Quit()
	b.out = b.Broadcast()
	b.fromListener = make(chan []byte, 1000)
	b.stop = make(chan bool)
	b.auth = map[string]interface{}{
		"QueueName":           "",
		"AccessKey":           "",
		"AccessSecret":        "",
		"MaxNumberOfMessages": "10",
	}
}

func (b *FromSQS) stopListening() {
	log.Println("attempting to stop SQS reader")
	if b.listening {
		b.stop <- true
		b.listening = false
	}
}

// Run is the block's main loop. Here we listen on the different channels we set up.
func (b *FromSQS) Run() {
	var err error

	for {
		select {
		case msgI := <-b.inrule:
			for k, _ := range b.auth {
				b.auth[k], err = util.ParseString(msgI, k)
				if err != nil {
					b.Error(err)
					break
				}
			}

			b.stopListening()
			go b.listener()
		case <-b.quit:
			b.stopListening()
			return
		case msg := <-b.fromListener:
			var outMsg interface{}
			err := json.Unmarshal(msg, &outMsg)
			if err != nil {
				b.Error(err)
				continue
			}
			b.out <- outMsg
		case MsgChan := <-b.queryrule:
			// deal with a query request
			MsgChan <- b.auth
		}
	}
}
<commit_msg>added a not-great flag that protects panics in the SQS reader<commit_after>package library

import (
	"encoding/json"
	"log"

	"github.com/goamz/goamz/aws"
	"github.com/goamz/goamz/sqs"

	"strconv"
	"sync"

	"github.com/nytlabs/streamtools/st/blocks" // blocks
	"github.com/nytlabs/streamtools/st/util"
)

// lots of this code stolen brazenly from JP https://github.com/jprobinson
func (b *FromSQS) listener() {
	b.lock.Lock()
	accessKey, ok := b.auth["AccessKey"].(string)
	if !ok {
		b.Error("could not assert AccessKey to a string")
		return
	}
	accessSecret, ok := b.auth["AccessSecret"].(string)
	if !ok {
		b.Error("could not assert AccessSecret to a string")
		return
	}
	queueName, ok := b.auth["QueueName"].(string)
	if !ok {
		b.Error("could not assert queue name to a string")
		return
	}
	maxNstring, ok := b.auth["MaxNumberOfMessages"].(string)
	if !ok {
		b.Error("could not assert MaxNumberOfMessages to a string")
		return
	}
	auth := aws.Auth{AccessKey: accessKey, SecretKey: accessSecret}
	sqsClient := sqs.New(auth, aws.USEast)
	queue, err := sqsClient.GetQueue(queueName)
	if err != nil {
		b.Error(err)
		return
	}

	maxN, err := strconv.Atoi(maxNstring)
	if err != nil {
		b.Error(err)
		return
	}

	b.listening = true
	b.lock.Unlock()

	var resp *sqs.ReceiveMessageResponse
	for {
		select {
		case <-b.stop:
			log.Println("Exiting SQS read loop")
			return
		default:
			resp, err = queue.ReceiveMessage(maxN)
			if err != nil {
				b.Error(err)
			}
			if len(resp.Messages) == 0 {
				break
			}
			for _, m := range resp.Messages {
				select {
				case b.fromListener <- []byte(m.Body):
				default:
					log.Println("discarding messages")
					log.Println(len(b.fromListener))
					continue
				}

				if _, err = queue.DeleteMessage(&m); err != nil {
					b.Error(err)
				}
			}
		}
	}
}

// specify those channels we're going to use to communicate with streamtools
type FromSQS struct {
	blocks.Block
	queryrule chan blocks.MsgChan
	inrule    blocks.MsgChan
	out       blocks.MsgChan
	quit      blocks.MsgChan

	lock         sync.Mutex
	listening    bool
	fromListener chan []byte
	auth         map[string]interface{}
	stop         chan bool
}

// we need to build a simple factory so that streamtools can make new blocks of this kind
func NewFromSQS() blocks.BlockInterface {
	return &FromSQS{}
}

// Setup is called once before running the block. We build up the channels and specify what kind of block this is.
func (b *FromSQS) Setup() {
	b.Kind = "Queue I/O"
	b.Desc = "reads from Amazon's SQS, emitting each line of JSON as a separate message"
	b.inrule = b.InRoute("rule")
	b.queryrule = b.QueryRoute("rule")
	b.quit = b.Quit()
	b.out = b.Broadcast()
	b.fromListener = make(chan []byte, 1000)
	b.stop = make(chan bool)
	b.auth = map[string]interface{}{
		"QueueName":           "",
		"AccessKey":           "",
		"AccessSecret":        "",
		"MaxNumberOfMessages": "10",
	}
}

func (b *FromSQS) stopListening() {
	log.Println("attempting to stop SQS reader")
	log.Println(b.listening)
	if b.listening {
		log.Println("sending stop")
		b.stop <- true
		log.Println("sent stop")
		b.listening = false
	}
}

// Run is the block's main loop. Here we listen on the different channels we set up.
func (b *FromSQS) Run() {
	var err error

	for {
		select {
		case msgI := <-b.inrule:
			for k, _ := range b.auth {
				b.auth[k], err = util.ParseString(msgI, k)
				if err != nil {
					b.Error(err)
					break
				}
			}

			b.stopListening()
			go b.listener()
		case <-b.quit:
			b.stopListening()
			return
		case msg := <-b.fromListener:
			var outMsg interface{}
			err := json.Unmarshal(msg, &outMsg)
			if err != nil {
				b.Error(err)
				continue
			}
			b.out <- outMsg
		case MsgChan := <-b.queryrule:
			// deal with a query request
			MsgChan <- b.auth
		}
	}
}
<|endoftext|>


<commit_before>package stager_test

import (
	"time"

	Bbs "github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	. "github.com/cloudfoundry-incubator/stager/stager"
	"github.com/cloudfoundry/gunk/timeprovider"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Stage", func() {
	var stager Stager
	var bbs *Bbs.BBS

	BeforeEach(func() {
		bbs = Bbs.New(etcdRunner.Adapter(), timeprovider.NewTimeProvider())
		compilers := map[string]string{
			"penguin":     "penguin-compiler",
			"rabbit_hole": "rabbit-hole-compiler",
		}
		stager = New(bbs, compilers)
	})

	Context("when file the server is available", func() {
		BeforeEach(func() {
			_, _, err := bbs.MaintainFileServerPresence(10*time.Second, "http://file-server.com/", "abc123")
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("creates a Task with staging instructions", func() {
			modelChannel, _, _ := bbs.WatchForDesiredTask()

			err := stager.Stage(models.StagingRequestFromCC{
				AppId:                          "bunny",
				TaskId:                         "hop",
				AppBitsDownloadUri:             "http://example-uri.com/bunny",
				BuildArtifactsCacheDownloadUri: "http://example-uri.com/bunny-droppings",
				Stack:           "rabbit_hole",
				FileDescriptors: 17,
				MemoryMB:        256,
				DiskMB:          1024,
				Buildpacks: []models.Buildpack{
					{Name: "zfirst", Key: "zfirst-buildpack", Url: "first-buildpack-url"},
					{Name: "asecond", Key: "asecond-buildpack", Url: "second-buildpack-url"},
				},
				Environment: []models.EnvironmentVariable{
					{"VCAP_APPLICATION", "foo"},
					{"VCAP_SERVICES", "bar"},
				},
			})
			Ω(err).ShouldNot(HaveOccurred())

			var task *models.Task
			Eventually(modelChannel).Should(Receive(&task))

			Ω(task.Guid).To(Equal("bunny-hop"))
			Ω(task.Stack).To(Equal("rabbit_hole"))
			Ω(task.Log.Guid).To(Equal("bunny"))
			Ω(task.Log.SourceName).To(Equal("STG"))
			Ω(task.FileDescriptors).To(Equal(17))
			Ω(task.Log.Index).To(BeNil())

			expectedActions := []models.ExecutorAction{
				models.EmitProgressFor(
					models.ExecutorAction{
						models.DownloadAction{
							From:     "http://file-server.com/v1/static/rabbit-hole-compiler",
							To:       "/tmp/compiler",
							Extract:  true,
							CacheKey: "smelter-rabbit_hole",
						},
					},
					"",
					"",
					"Failed to Download Smelter",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.DownloadAction{
							From:    "http://example-uri.com/bunny",
							To:      "/app",
							Extract: true,
						},
					},
					"Downloading App Package",
					"Downloaded App Package",
					"Failed to Download App Package",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.DownloadAction{
							From:     "first-buildpack-url",
							To:       "/tmp/buildpacks/0fe7d5fc3f73b0ab8682a664da513fbd",
							Extract:  true,
							CacheKey: "zfirst-buildpack",
						},
					},
					"Downloading Buildpack: zfirst",
					"Downloaded Buildpack: zfirst",
					"Failed to Download Buildpack: zfirst",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.DownloadAction{
							From:     "second-buildpack-url",
							To:       "/tmp/buildpacks/58015c32d26f0ad3418f87dd9bf47797",
							Extract:  true,
							CacheKey: "asecond-buildpack",
						},
					},
					"Downloading Buildpack: asecond",
					"Downloaded Buildpack: asecond",
					"Failed to Download Buildpack: asecond",
				),
				models.Try(
					models.EmitProgressFor(
						models.ExecutorAction{
							models.DownloadAction{
								From:    "http://example-uri.com/bunny-droppings",
								To:      "/tmp/cache",
								Extract: true,
							},
						},
						"Downloading Build Artifacts Cache",
						"Downloaded Build Artifacts Cache",
						"No Build Artifacts Cache Found.  Proceeding...",
					),
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.RunAction{
							Script: "/tmp/compiler/run" +
								" -appDir='/app'" +
								" -buildArtifactsCacheDir='/tmp/cache'" +
								" -buildpackOrder='zfirst-buildpack,asecond-buildpack'" +
								" -buildpacksDir='/tmp/buildpacks'" +
								" -outputDir='/tmp/droplet'" +
								" -resultDir='/tmp/result'",
							Env: []models.EnvironmentVariable{
								{"VCAP_APPLICATION", "foo"},
								{"VCAP_SERVICES", "bar"},
							},
							Timeout: 15 * time.Minute,
						},
					},
					"Staging...",
					"Staging Complete",
					"Staging Failed",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.UploadAction{
							From:     "/tmp/droplet/",
							To:       "http://file-server.com/v1/droplet/bunny",
							Compress: false,
						},
					},
					"Uploading Droplet",
					"Droplet Uploaded",
					"Failed to Upload Droplet",
				),
				models.Try(
					models.EmitProgressFor(
						models.ExecutorAction{
							models.UploadAction{
								From:     "/tmp/cache/",
								To:       "http://file-server.com/v1/build_artifacts/bunny",
								Compress: true,
							},
						},
						"Uploading Build Artifacts Cache",
						"Uploaded Build Artifacts Cache",
						"Failed to Upload Build Artifacts Cache.  Proceeding...",
					),
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.FetchResultAction{
							File: "/tmp/result/result.json",
						},
					},
					"",
					"",
					"Failed to Fetch Detected Buildpack",
				),
			}

			for i, action := range task.Actions {
				Ω(action).To(Equal(expectedActions[i]))
			}

			Ω(task.MemoryMB).To(Equal(256))
			Ω(task.DiskMB).To(Equal(1024))
		})
	})

	Context("when build artifacts download uris are not provided", func() {
		BeforeEach(func() {
			_, _, err := bbs.MaintainFileServerPresence(10*time.Second, "http://file-server.com/", "abc123")
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("does not instruct the executor to download the cache", func() {
			modelChannel, _, _ := bbs.WatchForDesiredTask()

			err := stager.Stage(models.StagingRequestFromCC{
				AppId:              "bunny",
				TaskId:             "hop",
				AppBitsDownloadUri: "http://example-uri.com/bunny",
				Stack:              "rabbit_hole",
				FileDescriptors:    17,
				MemoryMB:           256,
				DiskMB:             1024,
				Buildpacks: []models.Buildpack{
					{Name: "zfirst", Key: "zfirst-buildpack", Url: "first-buildpack-url"},
					{Name: "asecond", Key: "asecond-buildpack", Url: "second-buildpack-url"},
				},
				Environment: []models.EnvironmentVariable{
					{"VCAP_APPLICATION", "foo"},
					{"VCAP_SERVICES", "bar"},
				},
			})
			Ω(err).ShouldNot(HaveOccurred())

			var task *models.Task
			Eventually(modelChannel).Should(Receive(&task))

			Ω(task.Guid).To(Equal("bunny-hop"))
			Ω(task.Stack).To(Equal("rabbit_hole"))
			Ω(task.Log.Guid).To(Equal("bunny"))
			Ω(task.Log.SourceName).To(Equal("STG"))
			Ω(task.FileDescriptors).To(Equal(17))
			Ω(task.Log.Index).To(BeNil())

			expectedActions := []models.ExecutorAction{
				models.EmitProgressFor(
					models.ExecutorAction{
						models.DownloadAction{
							From:     "http://file-server.com/v1/static/rabbit-hole-compiler",
							To:       "/tmp/compiler",
							Extract:  true,
							CacheKey: "smelter-rabbit_hole",
						},
					},
					"",
					"",
					"Failed to Download Smelter",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.DownloadAction{
							From:    "http://example-uri.com/bunny",
							To:      "/app",
							Extract: true,
						},
					},
					"Downloading App Package",
					"Downloaded App Package",
					"Failed to Download App Package",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.DownloadAction{
							From:     "first-buildpack-url",
							To:       "/tmp/buildpacks/0fe7d5fc3f73b0ab8682a664da513fbd",
							Extract:  true,
							CacheKey: "zfirst-buildpack",
						},
					},
					"Downloading Buildpack: zfirst",
					"Downloaded Buildpack: zfirst",
					"Failed to Download Buildpack: zfirst",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.DownloadAction{
							From:     "second-buildpack-url",
							To:       "/tmp/buildpacks/58015c32d26f0ad3418f87dd9bf47797",
							Extract:  true,
							CacheKey: "asecond-buildpack",
						},
					},
					"Downloading Buildpack: asecond",
					"Downloaded Buildpack: asecond",
					"Failed to Download Buildpack: asecond",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.RunAction{
							Script: "/tmp/compiler/run" +
								" -appDir='/app'" +
								" -buildArtifactsCacheDir='/tmp/cache'" +
								" -buildpackOrder='zfirst-buildpack,asecond-buildpack'" +
								" -buildpacksDir='/tmp/buildpacks'" +
								" -outputDir='/tmp/droplet'" +
								" -resultDir='/tmp/result'",
							Env: []models.EnvironmentVariable{
								{"VCAP_APPLICATION", "foo"},
								{"VCAP_SERVICES", "bar"},
							},
							Timeout: 15 * time.Minute,
						},
					},
					"Staging...",
					"Staging Complete",
					"Staging Failed",
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.UploadAction{
							From:     "/tmp/droplet/",
							To:       "http://file-server.com/v1/droplet/bunny",
							Compress: false,
						},
					},
					"Uploading Droplet",
					"Droplet Uploaded",
					"Failed to Upload Droplet",
				),
				models.Try(
					models.EmitProgressFor(
						models.ExecutorAction{
							models.UploadAction{
								From:     "/tmp/cache/",
								To:       "http://file-server.com/v1/build_artifacts/bunny",
								Compress: true,
							},
						},
						"Uploading Build Artifacts Cache",
						"Uploaded Build Artifacts Cache",
						"Failed to Upload Build Artifacts Cache.  Proceeding...",
					),
				),
				models.EmitProgressFor(
					models.ExecutorAction{
						models.FetchResultAction{
							File: "/tmp/result/result.json",
						},
					},
					"",
					"",
					"Failed to Fetch Detected Buildpack",
				),
			}

			for i, action := range task.Actions {
				Ω(action).To(Equal(expectedActions[i]))
			}

			Ω(task.MemoryMB).To(Equal(256))
			Ω(task.DiskMB).To(Equal(1024))
		})
	})

	Context("when build artifacts download url is not a valid url", func() {
		It("return a url parsing error", func() {
			err := stager.Stage(models.StagingRequestFromCC{
				AppId:                          "bunny",
				TaskId:                         "hop",
				AppBitsDownloadUri:             "http://example-uri.com/bunny",
				BuildArtifactsCacheDownloadUri: "not-a-url",
				Stack:           "rabbit_hole",
				FileDescriptors: 17,
				MemoryMB:        256,
				DiskMB:          1024,
				Buildpacks: []models.Buildpack{
					{Key: "zfirst-buildpack", Url: "first-buildpack-url"},
					{Key: "asecond-buildpack", Url: "second-buildpack-url"},
				},
				Environment: []models.EnvironmentVariable{
					{"VCAP_APPLICATION", "foo"},
					{"VCAP_SERVICES", "bar"},
				},
			})
			Ω(err).Should(HaveOccurred())
		})
	})

	Context("when file server is not available", func() {
		It("should return an error", func() {
			err := stager.Stage(models.StagingRequestFromCC{
				AppId:                          "bunny",
				TaskId:                         "hop",
				AppBitsDownloadUri:             "http://example-uri.com/bunny",
				BuildArtifactsCacheDownloadUri: "http://example-uri.com/bunny-droppings",
				Stack:    "rabbit_hole",
				MemoryMB: 256,
				DiskMB:   1024,
			})

			Ω(err).Should(HaveOccurred())
			Ω(err.Error()).Should(Equal("no available file server present"))
		})
	})

	Context("when no compiler is defined for the requested stack in stager configuration", func() {
		BeforeEach(func() {
			_, _, err := bbs.MaintainFileServerPresence(10*time.Second, "http://file-server.com/", "abc123")
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("should return an error", func() {
			bbs.WatchForDesiredTask()

			err := stager.Stage(models.StagingRequestFromCC{
				AppId:                          "bunny",
				TaskId:                         "hop",
				AppBitsDownloadUri:             "http://example-uri.com/bunny",
				BuildArtifactsCacheDownloadUri: "http://example-uri.com/bunny-droppings",
				Stack:    "no_such_stack",
				MemoryMB: 256,
				DiskMB:   1024,
			})

			Ω(err).Should(HaveOccurred())
			Ω(err.Error()).Should(Equal("no compiler defined for requested stack"))
		})
	})
})
<commit_msg>refactor stager tests<commit_after>package stager_test

import (
	"time"

	Bbs "github.com/cloudfoundry-incubator/runtime-schema/bbs"
	"github.com/cloudfoundry-incubator/runtime-schema/models"
	. "github.com/cloudfoundry-incubator/stager/stager"
	"github.com/cloudfoundry/gunk/timeprovider"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Stage", func() {
	var (
		stager                        Stager
		bbs                           *Bbs.BBS
		stagingRequest                models.StagingRequestFromCC
		downloadSmelterAction         models.ExecutorAction
		downloadAppAction             models.ExecutorAction
		downloadFirstBuildpackAction  models.ExecutorAction
		downloadSecontBuildpackAction models.ExecutorAction
		downloadBuildArtifactsAction  models.ExecutorAction
		runAction                     models.ExecutorAction
		uploadDropletAction           models.ExecutorAction
		uploadBuildArtifactsAction    models.ExecutorAction
		fetchResultsAction            models.ExecutorAction
	)

	BeforeEach(func() {
		bbs = Bbs.New(etcdRunner.Adapter(), timeprovider.NewTimeProvider())
		compilers := map[string]string{
			"penguin":     "penguin-compiler",
			"rabbit_hole": "rabbit-hole-compiler",
		}
		stager = New(bbs, compilers)

		stagingRequest = models.StagingRequestFromCC{
			AppId:                          "bunny",
			TaskId:                         "hop",
			AppBitsDownloadUri:             "http://example-uri.com/bunny",
			BuildArtifactsCacheDownloadUri: "http://example-uri.com/bunny-droppings",
			Stack:           "rabbit_hole",
			FileDescriptors: 17,
			MemoryMB:        256,
			DiskMB:          1024,
			Buildpacks: []models.Buildpack{
				{Name: "zfirst", Key: "zfirst-buildpack", Url: "first-buildpack-url"},
				{Name: "asecond", Key: "asecond-buildpack", Url: "second-buildpack-url"},
			},
			Environment: []models.EnvironmentVariable{
				{"VCAP_APPLICATION", "foo"},
				{"VCAP_SERVICES", "bar"},
			},
		}

		downloadSmelterAction = models.EmitProgressFor(
			models.ExecutorAction{
				models.DownloadAction{
					From:     "http://file-server.com/v1/static/rabbit-hole-compiler",
					To:       "/tmp/compiler",
					Extract:  true,
					CacheKey: "smelter-rabbit_hole",
				},
			},
			"",
			"",
			"Failed to Download Smelter",
		)

		downloadAppAction = models.EmitProgressFor(
			models.ExecutorAction{
				models.DownloadAction{
					From:    "http://example-uri.com/bunny",
					To:      "/app",
					Extract: true,
				},
			},
			"Downloading App Package",
			"Downloaded App Package",
			"Failed to Download App Package",
		)

		downloadFirstBuildpackAction = models.EmitProgressFor(
			models.ExecutorAction{
				models.DownloadAction{
					From:     "first-buildpack-url",
					To:       "/tmp/buildpacks/0fe7d5fc3f73b0ab8682a664da513fbd",
					Extract:  true,
					CacheKey: "zfirst-buildpack",
				},
			},
			"Downloading Buildpack: zfirst",
			"Downloaded Buildpack: zfirst",
			"Failed to Download Buildpack: zfirst",
		)

		downloadSecontBuildpackAction = models.EmitProgressFor(
			models.ExecutorAction{
				models.DownloadAction{
					From:     "second-buildpack-url",
					To:       "/tmp/buildpacks/58015c32d26f0ad3418f87dd9bf47797",
					Extract:  true,
					CacheKey: "asecond-buildpack",
				},
			},
			"Downloading Buildpack: asecond",
			"Downloaded Buildpack: asecond",
			"Failed to Download Buildpack: asecond",
		)

		downloadBuildArtifactsAction = models.Try(
			models.EmitProgressFor(
				models.ExecutorAction{
					models.DownloadAction{
						From:    "http://example-uri.com/bunny-droppings",
						To:      "/tmp/cache",
						Extract: true,
					},
				},
				"Downloading Build Artifacts Cache",
				"Downloaded Build Artifacts Cache",
				"No Build Artifacts Cache Found.  Proceeding...",
			),
		)

		runAction = models.EmitProgressFor(
			models.ExecutorAction{
				models.RunAction{
					Script: "/tmp/compiler/run" +
						" -appDir='/app'" +
						" -buildArtifactsCacheDir='/tmp/cache'" +
						" -buildpackOrder='zfirst-buildpack,asecond-buildpack'" +
						" -buildpacksDir='/tmp/buildpacks'" +
						" -outputDir='/tmp/droplet'" +
						" -resultDir='/tmp/result'",
					Env: []models.EnvironmentVariable{
						{"VCAP_APPLICATION", "foo"},
						{"VCAP_SERVICES", "bar"},
					},
					Timeout: 15 * time.Minute,
				},
			},
			"Staging...",
			"Staging Complete",
			"Staging Failed",
		)

		uploadDropletAction = models.EmitProgressFor(
			models.ExecutorAction{
				models.UploadAction{
					From:     "/tmp/droplet/",
					To:       "http://file-server.com/v1/droplet/bunny",
					Compress: false,
				},
			},
			"Uploading Droplet",
			"Droplet Uploaded",
			"Failed to Upload Droplet",
		)

		uploadBuildArtifactsAction = models.Try(
			models.EmitProgressFor(
				models.ExecutorAction{
					models.UploadAction{
						From:     "/tmp/cache/",
						To:       "http://file-server.com/v1/build_artifacts/bunny",
						Compress: true,
					},
				},
				"Uploading Build Artifacts Cache",
				"Uploaded Build Artifacts Cache",
				"Failed to Upload Build Artifacts Cache.  Proceeding...",
			),
		)

		fetchResultsAction = models.EmitProgressFor(
			models.ExecutorAction{
				models.FetchResultAction{
					File: "/tmp/result/result.json",
				},
			},
			"",
			"",
			"Failed to Fetch Detected Buildpack",
		)
	})

	Context("when file the server is available", func() {
		BeforeEach(func() {
			_, _, err := bbs.MaintainFileServerPresence(10*time.Second, "http://file-server.com/", "abc123")
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("creates a Task with staging instructions", func() {
			modelChannel, _, _ := bbs.WatchForDesiredTask()

			err := stager.Stage(stagingRequest)
			Ω(err).ShouldNot(HaveOccurred())

			var task *models.Task
			Eventually(modelChannel).Should(Receive(&task))

			Ω(task.Guid).To(Equal("bunny-hop"))
			Ω(task.Stack).To(Equal("rabbit_hole"))
			Ω(task.Log.Guid).To(Equal("bunny"))
			Ω(task.Log.SourceName).To(Equal("STG"))
			Ω(task.FileDescriptors).To(Equal(17))
			Ω(task.Log.Index).To(BeNil())

			expectedActions := []models.ExecutorAction{
				downloadSmelterAction,
				downloadAppAction,
				downloadFirstBuildpackAction,
				downloadSecontBuildpackAction,
				downloadBuildArtifactsAction,
				runAction,
				uploadDropletAction,
				uploadBuildArtifactsAction,
				fetchResultsAction,
			}

			for i, action := range task.Actions {
				Ω(action).To(Equal(expectedActions[i]))
			}

			Ω(task.MemoryMB).To(Equal(256))
			Ω(task.DiskMB).To(Equal(1024))
		})
	})

	Context("when build artifacts download uris are not provided", func() {
		BeforeEach(func() {
			_, _, err := bbs.MaintainFileServerPresence(10*time.Second, "http://file-server.com/", "abc123")
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("does not instruct the executor to download the cache", func() {
			modelChannel, _, _ := bbs.WatchForDesiredTask()

			stagingRequest.BuildArtifactsCacheDownloadUri = ""
			err := stager.Stage(stagingRequest)
			Ω(err).ShouldNot(HaveOccurred())

			var task *models.Task
			Eventually(modelChannel).Should(Receive(&task))

			expectedActions := []models.ExecutorAction{
				downloadSmelterAction,
				downloadAppAction,
				downloadFirstBuildpackAction,
				downloadSecontBuildpackAction,
				runAction,
				uploadDropletAction,
				uploadBuildArtifactsAction,
				fetchResultsAction,
			}

			for i, action := range task.Actions {
				Ω(action).To(Equal(expectedActions[i]))
			}
		})
	})

	Context("when build artifacts download url is not a valid url", func() {
		It("return a url parsing error", func() {
			err := stager.Stage(models.StagingRequestFromCC{
				AppId:                          "bunny",
				TaskId:                         "hop",
				AppBitsDownloadUri:             "http://example-uri.com/bunny",
				BuildArtifactsCacheDownloadUri: "not-a-url",
				Stack:           "rabbit_hole",
				FileDescriptors: 17,
				MemoryMB:        256,
				DiskMB:          1024,
				Buildpacks: []models.Buildpack{
					{Key: "zfirst-buildpack", Url: "first-buildpack-url"},
					{Key: "asecond-buildpack", Url: "second-buildpack-url"},
				},
				Environment: []models.EnvironmentVariable{
					{"VCAP_APPLICATION", "foo"},
					{"VCAP_SERVICES", "bar"},
				},
			})
			Ω(err).Should(HaveOccurred())
		})
	})

	Context("when file server is not available", func() {
		It("should return an error", func() {
			err := stager.Stage(models.StagingRequestFromCC{
				AppId:                          "bunny",
				TaskId:                         "hop",
				AppBitsDownloadUri:             "http://example-uri.com/bunny",
				BuildArtifactsCacheDownloadUri: "http://example-uri.com/bunny-droppings",
				Stack:    "rabbit_hole",
				MemoryMB: 256,
				DiskMB:   1024,
			})

			Ω(err).Should(HaveOccurred())
			Ω(err.Error()).Should(Equal("no available file server present"))
		})
	})

	Context("when no compiler is defined for the requested stack in stager configuration", func() {
		BeforeEach(func() {
			_, _, err := bbs.MaintainFileServerPresence(10*time.Second, "http://file-server.com/", "abc123")
			Ω(err).ShouldNot(HaveOccurred())
		})

		It("should return an error", func() {
			bbs.WatchForDesiredTask()

			err := stager.Stage(models.StagingRequestFromCC{
				AppId:                          "bunny",
				TaskId:                         "hop",
				AppBitsDownloadUri:             "http://example-uri.com/bunny",
				BuildArtifactsCacheDownloadUri: "http://example-uri.com/bunny-droppings",
				Stack:    "no_such_stack",
				MemoryMB: 256,
				DiskMB:   1024,
			})

			Ω(err).Should(HaveOccurred())
			Ω(err.Error()).Should(Equal("no compiler defined for requested stack"))
		})
	})
})
<|endoftext|>


<commit_before>package state

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/logger"
	"github.com/ninjasphere/go-ninja/model"
)

// subcribe to all $device/+/channel/+/event/state

// store in a hash by $device/:deviceid/channel/:channelid/event/:eventid
//

// struct date, payload
type LastState struct {
	Timestamp int64       `json:"timestamp"`
	Payload   interface{} `json:"payload"`
}

// merge state

type StateManager interface {
	Merge(thing *model.Thing)
	Reset()
}

type NinjaStateManager struct {
	sync.Mutex
	log        *logger.Logger
	conn       *ninja.Connection
	lastStates map[string]*LastState
}

func NewStateManager(conn *ninja.Connection) StateManager {

	sm := &NinjaStateManager{
		lastStates: make(map[string]*LastState),
		conn:       conn,
		log:        logger.GetLogger("sphere-go-homecloud-state"),
	}

	go sm.startListener()

	return sm
}

func (sm *NinjaStateManager) Merge(thing *model.Thing) {

	deviceID := thing.DeviceID

	if thing.Device != nil && thing.Device.Channels != nil {
		for _, channelModel := range *thing.Device.Channels {

			key := fmt.Sprintf("%s-%s", *deviceID, channelModel.ID)

			sm.log.Infof("channel key %s state %v", key, sm.lastStates[key])

			if sm.lastStates[key] != nil {
				channelModel.LastState = sm.lastStates[key]
			}

		}
	}

}

func (sm *NinjaStateManager) Reset() {
	defer sm.Unlock()
	sm.Lock()
	sm.lastStates = make(map[string]*LastState)
}

func (sm *NinjaStateManager) startListener() {

	sm.log.Infof("startListener")

	err := sm.conn.GetServiceClient("$device/:deviceid/channel/:channelid").OnEvent("state", func(params *json.RawMessage, values map[string]string) bool {

		var data interface{}

		err := json.Unmarshal(*params, &data)

		if err != nil {
			sm.log.Errorf("bad content: %s", err)
		}

		key := fmt.Sprintf("%s-%s", values["deviceid"], values["channelid"])

		sm.lastStates[key] = &LastState{
			Timestamp: int64(time.Now().UnixNano() / 1e6),
			Payload:   data,
		}

		return true
	})

	if err != nil {
		sm.log.FatalError(err, "cant register service")
	}

}
<commit_msg>Check for key before setting.<commit_after>package state

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"github.com/ninjasphere/go-ninja/api"
	"github.com/ninjasphere/go-ninja/logger"
	"github.com/ninjasphere/go-ninja/model"
)

// struct date, payload
type LastState struct {
	Timestamp int64       `json:"timestamp"`
	Payload   interface{} `json:"payload"`
}

// merge state

type StateManager interface {
	Merge(thing *model.Thing)
	Reset()
}

type NinjaStateManager struct {
	sync.Mutex
	log        *logger.Logger
	conn       *ninja.Connection
	lastStates map[string]*LastState
}

func NewStateManager(conn *ninja.Connection) StateManager {

	sm := &NinjaStateManager{
		lastStates: make(map[string]*LastState),
		conn:       conn,
		log:        logger.GetLogger("sphere-go-homecloud-state"),
	}

	go sm.startListener()

	return sm
}

func (sm *NinjaStateManager) Merge(thing *model.Thing) {

	deviceID := thing.DeviceID

	if thing.Device != nil && thing.Device.Channels != nil {
		for _, channelModel := range *thing.Device.Channels {

			key := fmt.Sprintf("%s-%s", *deviceID, channelModel.ID)

			sm.log.Infof("channel key %s state %v", key, sm.lastStates[key])

			if val, ok := sm.lastStates[key]; ok {
				channelModel.LastState = val
			}
		}
	}

}

func (sm *NinjaStateManager) Reset() {
	defer sm.Unlock()
	sm.Lock()
	sm.lastStates = make(map[string]*LastState)
}

func (sm *NinjaStateManager) startListener() {

	sm.log.Infof("startListener")

	err := sm.conn.GetServiceClient("$device/:deviceid/channel/:channelid").OnEvent("state", func(params *json.RawMessage, values map[string]string) bool {

		var data interface{}

		err := json.Unmarshal(*params, &data)

		if err != nil {
			sm.log.Errorf("bad content: %s", err)
		}

		key := fmt.Sprintf("%s-%s", values["deviceid"], values["channelid"])

		sm.lastStates[key] = &LastState{
			Timestamp: int64(time.Now().UnixNano() / 1e6),
			Payload:   data,
		}

		return true
	})

	if err != nil {
		sm.log.FatalError(err, "cant register service")
	}

}
<|endoftext|>


<commit_before>package fake

import (
	"sync"
)

type FakeMetricSender struct {
	counters         map[string]uint64
	values           map[string]Metric
	containerMetrics map[string]ContainerMetric
	sync.RWMutex
}

type Metric struct {
	Value float64
	Unit  string
}

type ContainerMetric struct {
	ApplicationId string
	InstanceIndex int32
	CpuPercentage float64
	MemoryBytes   uint64
	DiskBytes     uint64
}

func NewFakeMetricSender() *FakeMetricSender {
	return &FakeMetricSender{
		counters:         make(map[string]uint64),
		values:           make(map[string]Metric),
		containerMetrics: make(map[string]ContainerMetric),
	}
}

func (fms *FakeMetricSender) SendValue(name string, value float64, unit string) error {
	fms.Lock()
	defer fms.Unlock()
	fms.values[name] = Metric{Value: value, Unit: unit}

	return nil
}

func (fms *FakeMetricSender) IncrementCounter(name string) error {
	fms.Lock()
	defer fms.Unlock()
	fms.counters[name]++

	return nil
}

func (fms *FakeMetricSender) AddToCounter(name string, delta uint64) error {
	fms.Lock()
	defer fms.Unlock()
	fms.counters[name] = fms.counters[name] + delta

	return nil
}

func (fms *FakeMetricSender) SendContainerMetric(applicationId string, instanceIndex int32, cpuPercentage float64, memoryBytes uint64, diskBytes uint64) error {
	fms.Lock()
	defer fms.Unlock()
	fms.containerMetrics[applicationId] = ContainerMetric{ApplicationId: applicationId, InstanceIndex: instanceIndex, CpuPercentage: cpuPercentage, MemoryBytes: memoryBytes, DiskBytes: diskBytes}

	return nil
}

func (fms *FakeMetricSender) GetValue(name string) Metric {
	fms.RLock()
	defer fms.RUnlock()

	return fms.values[name]
}

func (fms *FakeMetricSender) GetCounter(name string) uint64 {
	fms.RLock()
	defer fms.RUnlock()

	return fms.counters[name]
}

func (fms *FakeMetricSender) GetContainerMetric(applicationId string) ContainerMetric {
	fms.RLock()
	defer fms.RUnlock()

	return fms.containerMetrics[applicationId]
}
<commit_msg>Add FakeMetricSender.Reset()<commit_after>package fake

import (
	"sync"
)

type FakeMetricSender struct {
	counters         map[string]uint64
	values           map[string]Metric
	containerMetrics map[string]ContainerMetric
	sync.RWMutex
}

type Metric struct {
	Value float64
	Unit  string
}

type ContainerMetric struct {
	ApplicationId string
	InstanceIndex int32
	CpuPercentage float64
	MemoryBytes   uint64
	DiskBytes     uint64
}

func NewFakeMetricSender() *FakeMetricSender {
	return &FakeMetricSender{
		counters:         make(map[string]uint64),
		values:           make(map[string]Metric),
		containerMetrics: make(map[string]ContainerMetric),
	}
}

func (fms *FakeMetricSender) SendValue(name string, value float64, unit string) error {
	fms.Lock()
	defer fms.Unlock()
	fms.values[name] = Metric{Value: value, Unit: unit}

	return nil
}

func (fms *FakeMetricSender) IncrementCounter(name string) error {
	fms.Lock()
	defer fms.Unlock()
	fms.counters[name]++

	return nil
}

func (fms *FakeMetricSender) AddToCounter(name string, delta uint64) error {
	fms.Lock()
	defer fms.Unlock()
	fms.counters[name] = fms.counters[name] + delta

	return nil
}

func (fms *FakeMetricSender) SendContainerMetric(applicationId string, instanceIndex int32, cpuPercentage float64, memoryBytes uint64, diskBytes uint64) error {
	fms.Lock()
	defer fms.Unlock()
	fms.containerMetrics[applicationId] = ContainerMetric{ApplicationId: applicationId, InstanceIndex: instanceIndex, CpuPercentage: cpuPercentage, MemoryBytes: memoryBytes, DiskBytes: diskBytes}

	return nil
}

func (fms *FakeMetricSender) GetValue(name string) Metric {
	fms.RLock()
	defer fms.RUnlock()

	return fms.values[name]
}

func (fms *FakeMetricSender) GetCounter(name string) uint64 {
	fms.RLock()
	defer fms.RUnlock()

	return fms.counters[name]
}

func (fms *FakeMetricSender) GetContainerMetric(applicationId string) ContainerMetric {
	fms.RLock()
	defer fms.RUnlock()

	return fms.containerMetrics[applicationId]
}

func (fms *FakeMetricSender) Reset() {
	fms.Lock()
	defer fms.Unlock()

	fms.counters = make(map[string]uint64)
	fms.values = make(map[string]Metric)
	fms.containerMetrics = make(map[string]ContainerMetric)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Pinterest.com

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 bender

import (
	"math/rand"
	"time"
)

// ExponentialIntervalGenerator creates an IntervalGenerator that outputs
// exponentially distributed intervals. The resulting arrivals constitute
// a Poisson process. The rate parameter is the average queries per second
// for the generator, and corresponds to the reciprocal of the lambda parameter
// to an exponential distribution. In English, if you want to generate 30 QPS
// on average, pass 30 as the value of rate.
func ExponentialIntervalGenerator(rate float64) IntervalGenerator {
	rate = rate / float64(time.Second)
	return func(_ int64) int64 {
		return int64(rand.ExpFloat64() / rate)
	}
}

// RampedExponentialIntervalGenerator increases the average rate linearely
// from 0 to rate during the given ramp duration and keeps rate afterwards.
func RampedExponentialIntervalGenerator(rate float64, ramp time.Duration) IntervalGenerator {
	rate = rate / float64(time.Second)
	framp := float64(ramp)
	start := time.Now().UnixNano()
	return func(now int64) int64 {
		elapsed := float64(now - start)
		factor := elapsed / framp
		if factor > 1 {
			factor = 1
		} else if factor < 0.05 {
			// Do not let effective rate drop below 20% of the
			// final targe rate: Otherwise the tests does not
			// realy start as basicaly no request are generated.
			factor = 0.05
		}
		return int64(rand.ExpFloat64() / (factor * rate))
	}
}

// UniformIntervalGenerator creates and IntervalGenerator that outputs 1/rate every time it is
// called. Boring, right?
func UniformIntervalGenerator(rate float64) IntervalGenerator {
	irate := int64(rate / float64(time.Second))
	return func(_ int64) int64 {
		return irate
	}
}
<commit_msg>bender: fix comment<commit_after>/*
Copyright 2014 Pinterest.com

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 bender

import (
	"math/rand"
	"time"
)

// ExponentialIntervalGenerator creates an IntervalGenerator that outputs
// exponentially distributed intervals. The resulting arrivals constitute
// a Poisson process. The rate parameter is the average queries per second
// for the generator, and corresponds to the reciprocal of the lambda parameter
// to an exponential distribution. In English, if you want to generate 30 QPS
// on average, pass 30 as the value of rate.
func ExponentialIntervalGenerator(rate float64) IntervalGenerator {
	rate = rate / float64(time.Second)
	return func(_ int64) int64 {
		return int64(rand.ExpFloat64() / rate)
	}
}

// RampedExponentialIntervalGenerator increases the average rate linearely
// from 0 to rate during the given ramp duration and keeps rate afterwards.
func RampedExponentialIntervalGenerator(rate float64, ramp time.Duration) IntervalGenerator {
	rate = rate / float64(time.Second)
	framp := float64(ramp)
	start := time.Now().UnixNano()
	return func(now int64) int64 {
		elapsed := float64(now - start)
		factor := elapsed / framp
		if factor > 1 {
			factor = 1
		} else if factor < 0.05 {
			// Do not let effective rate drop below 5% of the
			// final targe rate: Otherwise the tests does not
			// really start as basicaly no request are generated.
			factor = 0.05
		}
		return int64(rand.ExpFloat64() / (factor * rate))
	}
}

// UniformIntervalGenerator creates and IntervalGenerator that outputs 1/rate every time it is
// called. Boring, right?
func UniformIntervalGenerator(rate float64) IntervalGenerator {
	irate := int64(rate / float64(time.Second))
	return func(_ int64) int64 {
		return irate
	}
}
<|endoftext|>


<commit_before>package protocol

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
)

// VersionNumber is a version number as int
type VersionNumber int32

// gQUIC version range as defined in the wiki: https://github.com/quicwg/base-drafts/wiki/QUIC-Versions
const (
	gquicVersion0   = 0x51303030
	maxGquicVersion = 0x51303439
)

// The version numbers, making grepping easier
const (
	Version39       VersionNumber = gquicVersion0 + 3*0x100 + 0x9 + iota
	VersionTLS      VersionNumber = 101
	VersionWhatever VersionNumber = 0 // for when the version doesn't matter
	VersionUnknown  VersionNumber = -1
)

// SupportedVersions lists the versions that the server supports
// must be in sorted descending order
var SupportedVersions = []VersionNumber{
	Version39,
}

// UsesTLS says if this QUIC version uses TLS 1.3 for the handshake
func (vn VersionNumber) UsesTLS() bool {
	return vn == VersionTLS
}

func (vn VersionNumber) String() string {
	switch vn {
	case VersionWhatever:
		return "whatever"
	case VersionUnknown:
		return "unknown"
	case VersionTLS:
		return "TLS dev version (WIP)"
	default:
		if vn.isGQUIC() {
			return fmt.Sprintf("gQUIC %d", vn.toGQUICVersion())
		}
		return fmt.Sprintf("%d", vn)
	}
}

// ToAltSvc returns the representation of the version for the H2 Alt-Svc parameters
func (vn VersionNumber) ToAltSvc() string {
	if vn.isGQUIC() {
		return fmt.Sprintf("%d", vn.toGQUICVersion())
	}
	return fmt.Sprintf("%d", vn)
}

// CryptoStreamID gets the Stream ID of the crypto stream
func (vn VersionNumber) CryptoStreamID() StreamID {
	if vn.isGQUIC() {
		return 1
	}
	return 0
}

// UsesIETFFrameFormat tells if this version uses the IETF frame format
func (vn VersionNumber) UsesIETFFrameFormat() bool {
	return vn != Version39
}

// StreamContributesToConnectionFlowControl says if a stream contributes to connection-level flow control
func (vn VersionNumber) StreamContributesToConnectionFlowControl(id StreamID) bool {
	if id == vn.CryptoStreamID() {
		return false
	}
	if vn.isGQUIC() && id == 3 {
		return false
	}
	return true
}

func (vn VersionNumber) isGQUIC() bool {
	return vn > gquicVersion0 && vn <= maxGquicVersion
}

func (vn VersionNumber) toGQUICVersion() int {
	return int(10*(vn-gquicVersion0)/0x100) + int(vn%0x10)
}

// IsSupportedVersion returns true if the server supports this version
func IsSupportedVersion(supported []VersionNumber, v VersionNumber) bool {
	for _, t := range supported {
		if t == v {
			return true
		}
	}
	return false
}

// ChooseSupportedVersion finds the best version in the overlap of ours and theirs
// ours is a slice of versions that we support, sorted by our preference (descending)
// theirs is a slice of versions offered by the peer. The order does not matter.
// The bool returned indicates if a matching version was found.
func ChooseSupportedVersion(ours, theirs []VersionNumber) (VersionNumber, bool) {
	for _, ourVer := range ours {
		for _, theirVer := range theirs {
			if ourVer == theirVer {
				return ourVer, true
			}
		}
	}
	return 0, false
}

// generateReservedVersion generates a reserved version number (v & 0x0f0f0f0f == 0x0a0a0a0a)
func generateReservedVersion() VersionNumber {
	b := make([]byte, 4)
	_, _ = rand.Read(b) // ignore the error here. Failure to read random data doesn't break anything
	return VersionNumber((binary.BigEndian.Uint32(b) | 0x0a0a0a0a) & 0xfafafafa)
}

// GetGreasedVersions adds one reserved version number to a slice of version numbers, at a random position
func GetGreasedVersions(supported []VersionNumber) []VersionNumber {
	b := make([]byte, 1)
	_, _ = rand.Read(b) // ignore the error here. Failure to read random data doesn't break anything
	randPos := int(b[0]) % (len(supported) + 1)
	greased := make([]VersionNumber, len(supported)+1)
	copy(greased, supported[:randPos])
	greased[randPos] = generateReservedVersion()
	copy(greased[randPos+1:], supported[randPos:])
	return greased
}
<commit_msg>make the version number a uint32<commit_after>package protocol

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"math"
)

// VersionNumber is a version number as int
type VersionNumber uint32

// gQUIC version range as defined in the wiki: https://github.com/quicwg/base-drafts/wiki/QUIC-Versions
const (
	gquicVersion0   = 0x51303030
	maxGquicVersion = 0x51303439
)

// The version numbers, making grepping easier
const (
	Version39       VersionNumber = gquicVersion0 + 3*0x100 + 0x9 + iota
	VersionTLS      VersionNumber = 101
	VersionWhatever VersionNumber = 0 // for when the version doesn't matter
	VersionUnknown  VersionNumber = math.MaxUint32
)

// SupportedVersions lists the versions that the server supports
// must be in sorted descending order
var SupportedVersions = []VersionNumber{
	Version39,
}

// UsesTLS says if this QUIC version uses TLS 1.3 for the handshake
func (vn VersionNumber) UsesTLS() bool {
	return vn == VersionTLS
}

func (vn VersionNumber) String() string {
	switch vn {
	case VersionWhatever:
		return "whatever"
	case VersionUnknown:
		return "unknown"
	case VersionTLS:
		return "TLS dev version (WIP)"
	default:
		if vn.isGQUIC() {
			return fmt.Sprintf("gQUIC %d", vn.toGQUICVersion())
		}
		return fmt.Sprintf("%d", vn)
	}
}

// ToAltSvc returns the representation of the version for the H2 Alt-Svc parameters
func (vn VersionNumber) ToAltSvc() string {
	if vn.isGQUIC() {
		return fmt.Sprintf("%d", vn.toGQUICVersion())
	}
	return fmt.Sprintf("%d", vn)
}

// CryptoStreamID gets the Stream ID of the crypto stream
func (vn VersionNumber) CryptoStreamID() StreamID {
	if vn.isGQUIC() {
		return 1
	}
	return 0
}

// UsesIETFFrameFormat tells if this version uses the IETF frame format
func (vn VersionNumber) UsesIETFFrameFormat() bool {
	return vn != Version39
}

// StreamContributesToConnectionFlowControl says if a stream contributes to connection-level flow control
func (vn VersionNumber) StreamContributesToConnectionFlowControl(id StreamID) bool {
	if id == vn.CryptoStreamID() {
		return false
	}
	if vn.isGQUIC() && id == 3 {
		return false
	}
	return true
}

func (vn VersionNumber) isGQUIC() bool {
	return vn > gquicVersion0 && vn <= maxGquicVersion
}

func (vn VersionNumber) toGQUICVersion() int {
	return int(10*(vn-gquicVersion0)/0x100) + int(vn%0x10)
}

// IsSupportedVersion returns true if the server supports this version
func IsSupportedVersion(supported []VersionNumber, v VersionNumber) bool {
	for _, t := range supported {
		if t == v {
			return true
		}
	}
	return false
}

// ChooseSupportedVersion finds the best version in the overlap of ours and theirs
// ours is a slice of versions that we support, sorted by our preference (descending)
// theirs is a slice of versions offered by the peer. The order does not matter.
// The bool returned indicates if a matching version was found.
func ChooseSupportedVersion(ours, theirs []VersionNumber) (VersionNumber, bool) {
	for _, ourVer := range ours {
		for _, theirVer := range theirs {
			if ourVer == theirVer {
				return ourVer, true
			}
		}
	}
	return 0, false
}

// generateReservedVersion generates a reserved version number (v & 0x0f0f0f0f == 0x0a0a0a0a)
func generateReservedVersion() VersionNumber {
	b := make([]byte, 4)
	_, _ = rand.Read(b) // ignore the error here. Failure to read random data doesn't break anything
	return VersionNumber((binary.BigEndian.Uint32(b) | 0x0a0a0a0a) & 0xfafafafa)
}

// GetGreasedVersions adds one reserved version number to a slice of version numbers, at a random position
func GetGreasedVersions(supported []VersionNumber) []VersionNumber {
	b := make([]byte, 1)
	_, _ = rand.Read(b) // ignore the error here. Failure to read random data doesn't break anything
	randPos := int(b[0]) % (len(supported) + 1)
	greased := make([]VersionNumber, len(supported)+1)
	copy(greased, supported[:randPos])
	greased[randPos] = generateReservedVersion()
	copy(greased[randPos+1:], supported[randPos:])
	return greased
}
<|endoftext|>


<commit_before>// Copyright 2016 The Ebiten 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 restorable

import (
	"errors"
	"image"
	"image/color"
	"runtime"

	"github.com/hajimehoshi/ebiten/internal/affine"
	"github.com/hajimehoshi/ebiten/internal/graphics"
	"github.com/hajimehoshi/ebiten/internal/opengl"
)

const drawImageHistoryMax = 100

type drawImageHistoryItem struct {
	image    *Image
	vertices []float32
	colorm   affine.ColorM
	mode     opengl.CompositeMode
}

// Image represents an image that can be restored when GL context is lost.
type Image struct {
	image  *graphics.Image
	filter opengl.Filter

	// baseImage and baseColor are exclusive.
	basePixels       []uint8
	baseColor        color.RGBA
	drawImageHistory []*drawImageHistoryItem
	stale            bool

	volatile bool
	screen   bool
}

func NewImage(width, height int, filter opengl.Filter, volatile bool) *Image {
	i := &Image{
		image:    graphics.NewImage(width, height, filter),
		filter:   filter,
		volatile: volatile,
	}
	theImages.add(i)
	runtime.SetFinalizer(i, (*Image).Dispose)
	return i
}

func NewImageFromImage(source *image.RGBA, width, height int, filter opengl.Filter) *Image {
	w2, h2 := graphics.NextPowerOf2Int(width), graphics.NextPowerOf2Int(height)
	p := make([]uint8, 4*w2*h2)
	for j := 0; j < height; j++ {
		copy(p[j*w2*4:(j+1)*w2*4], source.Pix[j*source.Stride:])
	}
	i := &Image{
		image:      graphics.NewImageFromImage(source, width, height, filter),
		basePixels: p,
		filter:     filter,
	}
	theImages.add(i)
	runtime.SetFinalizer(i, (*Image).Dispose)
	return i
}

func NewScreenFramebufferImage(width, height int) *Image {
	i := &Image{
		image:    graphics.NewScreenFramebufferImage(width, height),
		volatile: true,
		screen:   true,
	}
	theImages.add(i)
	runtime.SetFinalizer(i, (*Image).Dispose)
	return i
}

func (p *Image) Size() (int, int) {
	return p.image.Size()
}

func (p *Image) makeStale() {
	p.basePixels = nil
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = true
}

func (p *Image) clearIfVolatile() {
	if !p.volatile {
		return
	}
	p.basePixels = nil
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
	if p.image == nil {
		panic("not reach")
	}
	p.image.Fill(color.RGBA{})
}

func (p *Image) Fill(clr color.RGBA) {
	theImages.resetPixelsIfDependingOn(p)
	p.basePixels = nil
	p.baseColor = clr
	p.drawImageHistory = nil
	p.stale = false
	p.image.Fill(clr)
}

func (p *Image) ReplacePixels(pixels []uint8) {
	theImages.resetPixelsIfDependingOn(p)
	p.image.ReplacePixels(pixels)
	p.basePixels = pixels
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
}

func (p *Image) DrawImage(img *Image, vertices []float32, colorm *affine.ColorM, mode opengl.CompositeMode) {
	theImages.resetPixelsIfDependingOn(p)
	if img.stale || img.volatile {
		p.makeStale()
	} else {
		p.appendDrawImageHistory(img, vertices, colorm, mode)
	}
	p.image.DrawImage(img.image, vertices, colorm, mode)
}

func (p *Image) appendDrawImageHistory(image *Image, vertices []float32, colorm *affine.ColorM, mode opengl.CompositeMode) {
	if p.stale {
		return
	}
	if len(p.drawImageHistory)+1 > drawImageHistoryMax {
		p.makeStale()
		return
	}
	// All images must be resolved and not stale each after frame.
	// So we don't have to care if image is stale or not here.
	item := &drawImageHistoryItem{
		image:    image,
		vertices: vertices,
		colorm:   *colorm,
		mode:     mode,
	}
	p.drawImageHistory = append(p.drawImageHistory, item)
}

// At returns a color value at (x, y).
//
// Note that this must not be called until context is available.
// This means Pixels members must match with acutal state in VRAM.
func (p *Image) At(x, y int, context *opengl.Context) (color.RGBA, error) {
	w, h := p.image.Size()
	w2, h2 := graphics.NextPowerOf2Int(w), graphics.NextPowerOf2Int(h)
	if x < 0 || y < 0 || w2 <= x || h2 <= y {
		return color.RGBA{}, nil
	}
	if p.basePixels == nil || p.drawImageHistory != nil || p.stale {
		if err := p.readPixelsFromVRAM(p.image, context); err != nil {
			return color.RGBA{}, err
		}
	}
	idx := 4*x + 4*y*w2
	r, g, b, a := p.basePixels[idx], p.basePixels[idx+1], p.basePixels[idx+2], p.basePixels[idx+3]
	return color.RGBA{r, g, b, a}, nil
}

func (p *Image) makeStaleIfDependingOn(target *Image) {
	if p.stale {
		return
	}
	for _, c := range p.drawImageHistory {
		if c.image == target {
			p.makeStale()
			return
		}
	}
}

func (p *Image) readPixelsFromVRAM(image *graphics.Image, context *opengl.Context) error {
	var err error
	p.basePixels, err = image.Pixels(context)
	if err != nil {
		return err
	}
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
	return nil
}

func (p *Image) resolveStalePixels(context *opengl.Context) error {
	if p.volatile {
		return nil
	}
	if !p.stale {
		return nil
	}
	return p.readPixelsFromVRAM(p.image, context)
}

func (p *Image) hasDependency() bool {
	if p.stale {
		return false
	}
	return len(p.drawImageHistory) > 0
}

// Restore restores *graphics.Image from the pixels using its state.
func (p *Image) restore(context *opengl.Context) error {
	w, h := p.image.Size()
	if p.screen {
		// The screen image should also be recreated because framebuffer might
		// be changed.
		p.image = graphics.NewScreenFramebufferImage(w, h)
		p.basePixels = nil
		p.baseColor = color.RGBA{}
		p.drawImageHistory = nil
		p.stale = false
		return nil
	}
	if p.volatile {
		p.image = graphics.NewImage(w, h, p.filter)
		p.basePixels = nil
		p.baseColor = color.RGBA{}
		p.drawImageHistory = nil
		p.stale = false
		return nil
	}
	if p.stale {
		// TODO: panic here?
		return errors.New("restorable: pixels must not be stale when restoring")
	}
	w2, h2 := graphics.NextPowerOf2Int(w), graphics.NextPowerOf2Int(h)
	img := image.NewRGBA(image.Rect(0, 0, w2, h2))
	if p.basePixels != nil {
		for j := 0; j < h; j++ {
			copy(img.Pix[j*img.Stride:], p.basePixels[j*w2*4:(j+1)*w2*4])
		}
	}
	gimg := graphics.NewImageFromImage(img, w, h, p.filter)
	if p.baseColor != (color.RGBA{}) {
		if p.basePixels != nil {
			panic("not reach")
		}
		gimg.Fill(p.baseColor)
	}
	for _, c := range p.drawImageHistory {
		// c.image.image must be already restored.
		if c.image.hasDependency() {
			panic("not reach")
		}
		gimg.DrawImage(c.image.image, c.vertices, &c.colorm, c.mode)
	}
	p.image = gimg

	var err error
	p.basePixels, err = gimg.Pixels(context)
	if err != nil {
		return err
	}
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
	return nil
}

func (p *Image) Dispose() {
	theImages.resetPixelsIfDependingOn(p)
	p.image.Dispose()
	p.image = nil
	p.basePixels = nil
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
	theImages.remove(p)
	runtime.SetFinalizer(p, nil)
}

func (p *Image) IsInvalidated(context *opengl.Context) bool {
	return p.image.IsInvalidated(context)
}
<commit_msg>restorable: VRAM -> GPU<commit_after>// Copyright 2016 The Ebiten 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 restorable

import (
	"errors"
	"image"
	"image/color"
	"runtime"

	"github.com/hajimehoshi/ebiten/internal/affine"
	"github.com/hajimehoshi/ebiten/internal/graphics"
	"github.com/hajimehoshi/ebiten/internal/opengl"
)

const drawImageHistoryMax = 100

type drawImageHistoryItem struct {
	image    *Image
	vertices []float32
	colorm   affine.ColorM
	mode     opengl.CompositeMode
}

// Image represents an image that can be restored when GL context is lost.
type Image struct {
	image  *graphics.Image
	filter opengl.Filter

	// baseImage and baseColor are exclusive.
	basePixels       []uint8
	baseColor        color.RGBA
	drawImageHistory []*drawImageHistoryItem
	stale            bool

	volatile bool
	screen   bool
}

func NewImage(width, height int, filter opengl.Filter, volatile bool) *Image {
	i := &Image{
		image:    graphics.NewImage(width, height, filter),
		filter:   filter,
		volatile: volatile,
	}
	theImages.add(i)
	runtime.SetFinalizer(i, (*Image).Dispose)
	return i
}

func NewImageFromImage(source *image.RGBA, width, height int, filter opengl.Filter) *Image {
	w2, h2 := graphics.NextPowerOf2Int(width), graphics.NextPowerOf2Int(height)
	p := make([]uint8, 4*w2*h2)
	for j := 0; j < height; j++ {
		copy(p[j*w2*4:(j+1)*w2*4], source.Pix[j*source.Stride:])
	}
	i := &Image{
		image:      graphics.NewImageFromImage(source, width, height, filter),
		basePixels: p,
		filter:     filter,
	}
	theImages.add(i)
	runtime.SetFinalizer(i, (*Image).Dispose)
	return i
}

func NewScreenFramebufferImage(width, height int) *Image {
	i := &Image{
		image:    graphics.NewScreenFramebufferImage(width, height),
		volatile: true,
		screen:   true,
	}
	theImages.add(i)
	runtime.SetFinalizer(i, (*Image).Dispose)
	return i
}

func (p *Image) Size() (int, int) {
	return p.image.Size()
}

func (p *Image) makeStale() {
	p.basePixels = nil
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = true
}

func (p *Image) clearIfVolatile() {
	if !p.volatile {
		return
	}
	p.basePixels = nil
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
	if p.image == nil {
		panic("not reach")
	}
	p.image.Fill(color.RGBA{})
}

func (p *Image) Fill(clr color.RGBA) {
	theImages.resetPixelsIfDependingOn(p)
	p.basePixels = nil
	p.baseColor = clr
	p.drawImageHistory = nil
	p.stale = false
	p.image.Fill(clr)
}

func (p *Image) ReplacePixels(pixels []uint8) {
	theImages.resetPixelsIfDependingOn(p)
	p.image.ReplacePixels(pixels)
	p.basePixels = pixels
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
}

func (p *Image) DrawImage(img *Image, vertices []float32, colorm *affine.ColorM, mode opengl.CompositeMode) {
	theImages.resetPixelsIfDependingOn(p)
	if img.stale || img.volatile {
		p.makeStale()
	} else {
		p.appendDrawImageHistory(img, vertices, colorm, mode)
	}
	p.image.DrawImage(img.image, vertices, colorm, mode)
}

func (p *Image) appendDrawImageHistory(image *Image, vertices []float32, colorm *affine.ColorM, mode opengl.CompositeMode) {
	if p.stale {
		return
	}
	if len(p.drawImageHistory)+1 > drawImageHistoryMax {
		p.makeStale()
		return
	}
	// All images must be resolved and not stale each after frame.
	// So we don't have to care if image is stale or not here.
	item := &drawImageHistoryItem{
		image:    image,
		vertices: vertices,
		colorm:   *colorm,
		mode:     mode,
	}
	p.drawImageHistory = append(p.drawImageHistory, item)
}

// At returns a color value at (x, y).
//
// Note that this must not be called until context is available.
// This means Pixels members must match with acutal state in GPU.
func (p *Image) At(x, y int, context *opengl.Context) (color.RGBA, error) {
	w, h := p.image.Size()
	w2, h2 := graphics.NextPowerOf2Int(w), graphics.NextPowerOf2Int(h)
	if x < 0 || y < 0 || w2 <= x || h2 <= y {
		return color.RGBA{}, nil
	}
	if p.basePixels == nil || p.drawImageHistory != nil || p.stale {
		if err := p.readPixelsFromGPU(p.image, context); err != nil {
			return color.RGBA{}, err
		}
	}
	idx := 4*x + 4*y*w2
	r, g, b, a := p.basePixels[idx], p.basePixels[idx+1], p.basePixels[idx+2], p.basePixels[idx+3]
	return color.RGBA{r, g, b, a}, nil
}

func (p *Image) makeStaleIfDependingOn(target *Image) {
	if p.stale {
		return
	}
	for _, c := range p.drawImageHistory {
		if c.image == target {
			p.makeStale()
			return
		}
	}
}

func (p *Image) readPixelsFromGPU(image *graphics.Image, context *opengl.Context) error {
	var err error
	p.basePixels, err = image.Pixels(context)
	if err != nil {
		return err
	}
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
	return nil
}

func (p *Image) resolveStalePixels(context *opengl.Context) error {
	if p.volatile {
		return nil
	}
	if !p.stale {
		return nil
	}
	return p.readPixelsFromGPU(p.image, context)
}

func (p *Image) hasDependency() bool {
	if p.stale {
		return false
	}
	return len(p.drawImageHistory) > 0
}

// Restore restores *graphics.Image from the pixels using its state.
func (p *Image) restore(context *opengl.Context) error {
	w, h := p.image.Size()
	if p.screen {
		// The screen image should also be recreated because framebuffer might
		// be changed.
		p.image = graphics.NewScreenFramebufferImage(w, h)
		p.basePixels = nil
		p.baseColor = color.RGBA{}
		p.drawImageHistory = nil
		p.stale = false
		return nil
	}
	if p.volatile {
		p.image = graphics.NewImage(w, h, p.filter)
		p.basePixels = nil
		p.baseColor = color.RGBA{}
		p.drawImageHistory = nil
		p.stale = false
		return nil
	}
	if p.stale {
		// TODO: panic here?
		return errors.New("restorable: pixels must not be stale when restoring")
	}
	w2, h2 := graphics.NextPowerOf2Int(w), graphics.NextPowerOf2Int(h)
	img := image.NewRGBA(image.Rect(0, 0, w2, h2))
	if p.basePixels != nil {
		for j := 0; j < h; j++ {
			copy(img.Pix[j*img.Stride:], p.basePixels[j*w2*4:(j+1)*w2*4])
		}
	}
	gimg := graphics.NewImageFromImage(img, w, h, p.filter)
	if p.baseColor != (color.RGBA{}) {
		if p.basePixels != nil {
			panic("not reach")
		}
		gimg.Fill(p.baseColor)
	}
	for _, c := range p.drawImageHistory {
		// c.image.image must be already restored.
		if c.image.hasDependency() {
			panic("not reach")
		}
		gimg.DrawImage(c.image.image, c.vertices, &c.colorm, c.mode)
	}
	p.image = gimg

	var err error
	p.basePixels, err = gimg.Pixels(context)
	if err != nil {
		return err
	}
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
	return nil
}

func (p *Image) Dispose() {
	theImages.resetPixelsIfDependingOn(p)
	p.image.Dispose()
	p.image = nil
	p.basePixels = nil
	p.baseColor = color.RGBA{}
	p.drawImageHistory = nil
	p.stale = false
	theImages.remove(p)
	runtime.SetFinalizer(p, nil)
}

func (p *Image) IsInvalidated(context *opengl.Context) bool {
	return p.image.IsInvalidated(context)
}
<|endoftext|>


<commit_before>package mongodb

import (
	"container/heap"
	"github.com/name5566/leaf/log"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"sync"
)

// session
type Session struct {
	*mgo.Session
	ref   int
	index int
}

// session heap
type SessionHeap []*Session

func (h SessionHeap) Len() int {
	return len(h)
}

func (h SessionHeap) Less(i, j int) bool {
	return h[i].ref < h[j].ref
}

func (h SessionHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
	h[i].index = i
	h[j].index = j
}

func (h *SessionHeap) Push(s interface{}) {
	s.(*Session).index = len(*h)
	*h = append(*h, s.(*Session))
}

func (h *SessionHeap) Pop() interface{} {
	l := len(*h)
	s := (*h)[l-1]
	s.index = -1
	*h = (*h)[:l-1]
	return s
}

type DialContext struct {
	sync.Mutex
	sessions SessionHeap
}

// goroutine safe
func Dial(url string, sessionNum int) (*DialContext, error) {
	if sessionNum <= 0 {
		sessionNum = 100
		log.Release("invalid sessionNum, reset to %v", sessionNum)
	}

	s, err := mgo.Dial(url)
	if err != nil {
		return nil, err
	}

	c := new(DialContext)

	// sessions
	c.sessions = make(SessionHeap, sessionNum)
	c.sessions[0] = &Session{s, 0, 0}
	for i := 1; i < sessionNum; i++ {
		c.sessions[i] = &Session{s.New(), 0, i}
	}
	heap.Init(&c.sessions)

	return c, nil
}

// goroutine safe
func (c *DialContext) Close() {
	c.Lock()
	for _, s := range c.sessions {
		s.Close()
		if s.ref != 0 {
			log.Error("session ref = %v", s.ref)
		}
	}
	c.Unlock()
}

// goroutine safe
func (c *DialContext) Ref() *Session {
	c.Lock()
	s := c.sessions[0]
	if s.ref == 0 {
		s.Refresh()
	}
	s.ref++
	heap.Fix(&c.sessions, 0)
	c.Unlock()

	return s
}

// goroutine safe
func (c *DialContext) UnRef(s *Session) {
	c.Lock()
	s.ref--
	heap.Fix(&c.sessions, s.index)
	c.Unlock()
}

// goroutine safe
func (c *DialContext) EnsureCounter(db string, collection string, id string) error {
	s := c.Ref()
	defer c.UnRef(s)

	err := s.DB(db).C(collection).Insert(bson.M{
		"_id": id,
		"seq": 0,
	})
	if mgo.IsDup(err) {
		return nil
	} else {
		return err
	}
}

// goroutine safe
func (c *DialContext) NextSeq(db string, collection string, id string) (int, error) {
	s := c.Ref()
	defer c.UnRef(s)

	var res struct {
		Seq int
	}
	_, err := s.DB(db).C(collection).FindId(id).Apply(mgo.Change{
		Update:    bson.M{"$inc": bson.M{"seq": 1}},
		ReturnNew: true,
	}, &res)

	return res.Seq, err
}

// goroutine safe
func (c *DialContext) EnsureIndex(db string, collection string, key []string) error {
	s := c.Ref()
	defer c.UnRef(s)

	return s.DB(db).C(collection).EnsureIndex(mgo.Index{
		Key:    key,
		Unique: false,
		Sparse: true,
	})
}

// goroutine safe
func (c *DialContext) EnsureUniqueIndex(db string, collection string, key []string) error {
	s := c.Ref()
	defer c.UnRef(s)

	return s.DB(db).C(collection).EnsureIndex(mgo.Index{
		Key:    key,
		Unique: true,
		Sparse: true,
	})
}
<commit_msg>add a function: DialWithTimeout<commit_after>package mongodb

import (
	"container/heap"
	"github.com/name5566/leaf/log"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"sync"
	"time"
)

// session
type Session struct {
	*mgo.Session
	ref   int
	index int
}

// session heap
type SessionHeap []*Session

func (h SessionHeap) Len() int {
	return len(h)
}

func (h SessionHeap) Less(i, j int) bool {
	return h[i].ref < h[j].ref
}

func (h SessionHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
	h[i].index = i
	h[j].index = j
}

func (h *SessionHeap) Push(s interface{}) {
	s.(*Session).index = len(*h)
	*h = append(*h, s.(*Session))
}

func (h *SessionHeap) Pop() interface{} {
	l := len(*h)
	s := (*h)[l-1]
	s.index = -1
	*h = (*h)[:l-1]
	return s
}

type DialContext struct {
	sync.Mutex
	sessions SessionHeap
}

// goroutine safe
func Dial(url string, sessionNum int) (*DialContext, error) {
	c, err := DialWithTimeout(url, sessionNum, 10*time.Second, 5*time.Minute)
	return c, err
}

// goroutine safe
func DialWithTimeout(url string, sessionNum int, dialTimeout time.Duration, timeout time.Duration) (*DialContext, error) {
	if sessionNum <= 0 {
		sessionNum = 100
		log.Release("invalid sessionNum, reset to %v", sessionNum)
	}

	s, err := mgo.DialWithTimeout(url, dialTimeout)
	if err != nil {
		return nil, err
	}
	s.SetSyncTimeout(timeout)
	s.SetSocketTimeout(timeout)

	c := new(DialContext)

	// sessions
	c.sessions = make(SessionHeap, sessionNum)
	c.sessions[0] = &Session{s, 0, 0}
	for i := 1; i < sessionNum; i++ {
		c.sessions[i] = &Session{s.New(), 0, i}
	}
	heap.Init(&c.sessions)

	return c, nil
}

// goroutine safe
func (c *DialContext) Close() {
	c.Lock()
	for _, s := range c.sessions {
		s.Close()
		if s.ref != 0 {
			log.Error("session ref = %v", s.ref)
		}
	}
	c.Unlock()
}

// goroutine safe
func (c *DialContext) Ref() *Session {
	c.Lock()
	s := c.sessions[0]
	if s.ref == 0 {
		s.Refresh()
	}
	s.ref++
	heap.Fix(&c.sessions, 0)
	c.Unlock()

	return s
}

// goroutine safe
func (c *DialContext) UnRef(s *Session) {
	c.Lock()
	s.ref--
	heap.Fix(&c.sessions, s.index)
	c.Unlock()
}

// goroutine safe
func (c *DialContext) EnsureCounter(db string, collection string, id string) error {
	s := c.Ref()
	defer c.UnRef(s)

	err := s.DB(db).C(collection).Insert(bson.M{
		"_id": id,
		"seq": 0,
	})
	if mgo.IsDup(err) {
		return nil
	} else {
		return err
	}
}

// goroutine safe
func (c *DialContext) NextSeq(db string, collection string, id string) (int, error) {
	s := c.Ref()
	defer c.UnRef(s)

	var res struct {
		Seq int
	}
	_, err := s.DB(db).C(collection).FindId(id).Apply(mgo.Change{
		Update:    bson.M{"$inc": bson.M{"seq": 1}},
		ReturnNew: true,
	}, &res)

	return res.Seq, err
}

// goroutine safe
func (c *DialContext) EnsureIndex(db string, collection string, key []string) error {
	s := c.Ref()
	defer c.UnRef(s)

	return s.DB(db).C(collection).EnsureIndex(mgo.Index{
		Key:    key,
		Unique: false,
		Sparse: true,
	})
}

// goroutine safe
func (c *DialContext) EnsureUniqueIndex(db string, collection string, key []string) error {
	s := c.Ref()
	defer c.UnRef(s)

	return s.DB(db).C(collection).EnsureIndex(mgo.Index{
		Key:    key,
		Unique: true,
		Sparse: true,
	})
}
<|endoftext|>