<commit_before>package tsacmd

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	gclient "github.com/cloudfoundry-incubator/garden/client"
	gconn "github.com/cloudfoundry-incubator/garden/client/connection"
	"github.com/concourse/atc"
	"github.com/concourse/tsa"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/rata"
	"golang.org/x/crypto/ssh"
)

type registrarSSHServer struct {
	logger            lager.Logger
	atcEndpoint       *rata.RequestGenerator
	tokenGenerator    tsa.TokenGenerator
	heartbeatInterval time.Duration
	cprInterval       time.Duration
	forwardHost       string
	config            *ssh.ServerConfig
	httpClient        *http.Client
}

func (server *registrarSSHServer) Serve(listener net.Listener) {
	for {
		c, err := listener.Accept()
		if err != nil {
			if !strings.Contains(err.Error(), "use of closed network connection") {
				server.logger.Error("failed-to-accept", err)
			}

			return
		}

		logger := server.logger.Session("connection")

		conn, chans, reqs, err := ssh.NewServerConn(c, server.config)
		if err != nil {
			logger.Info("handshake-failed", lager.Data{"error": err.Error()})
			continue
		}

		go server.handleConn(logger, conn, chans, reqs)
	}
}

type forwardedTCPIP struct {
	bindAddr  string
	process   ifrit.Process
	boundPort uint32
}

func (server *registrarSSHServer) handleConn(logger lager.Logger, conn *ssh.ServerConn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request) {
	defer conn.Close()

	forwardedTCPIPs := make(chan forwardedTCPIP, 2)
	go server.handleForwardRequests(logger, conn, reqs, forwardedTCPIPs)

	var processes []ifrit.Process
	var process ifrit.Process

	// ensure processes get cleaned up
	defer func() {
		cleanupLog := logger.Session("cleanup")

		for _, p := range processes {
			cleanupLog.Debug("interrupting")

			p.Signal(os.Interrupt)
		}

		for _, p := range processes {
			err := <-p.Wait()
			if err != nil {
				cleanupLog.Error("process-exited-with-failure", err)
			} else {
				cleanupLog.Debug("process-exited-successfully")
			}
		}
	}()

	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			logger.Info("rejecting-unknown-channel-type", lager.Data{
				"type": newChannel.ChannelType(),
			})

			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		channel, requests, err := newChannel.Accept()
		if err != nil {
			logger.Error("failed-to-accept-channel", err)
			return
		}

		defer channel.Close()

		for req := range requests {
			logger.Info("channel-request", lager.Data{
				"type": req.Type,
			})

			if req.Type != "exec" {
				logger.Info("rejecting")
				req.Reply(false, nil)
				continue
			}

			var request execRequest
			err = ssh.Unmarshal(req.Payload, &request)
			if err != nil {
				logger.Error("malformed-exec-request", err)
				req.Reply(false, nil)
				return
			}

			workerRequest, err := parseRequest(request.Command)
			if err != nil {
				fmt.Fprintf(channel, "invalid command: %s", err)
				req.Reply(false, nil)
				continue
			}

			switch r := workerRequest.(type) {
			case registerWorkerRequest:
				logger := logger.Session("register-worker")

				req.Reply(true, nil)

				process, err = server.continuouslyRegisterWorkerDirectly(logger, channel)
				if err != nil {
					logger.Error("failed-to-register", err)
					return
				}

				processes = append(processes, process)

				err = conn.Wait()
				logger.Error("connection-closed", err)

			case forwardWorkerRequest:
				logger := logger.Session("forward-worker")

				req.Reply(true, nil)

				forwards := map[string]forwardedTCPIP{}

				for i := 0; i < r.expectedForwards(); i++ {
					select {
					case forwarded := <-forwardedTCPIPs:
						logger.Info("forwarded-tcpip", lager.Data{
							"bound-port": forwarded.boundPort,
						})

						processes = append(processes, forwarded.process)

						forwards[forwarded.bindAddr] = forwarded

					case <-time.After(10 * time.Second): // todo better?
						logger.Info("never-forwarded-tcpip")
					}
				}

				switch len(forwards) {
				case 0:
					fmt.Fprintf(channel, "requested forwarding but no forwards given\n")
					return

				case 1:
					for _, gardenForward := range forwards {
						process, err = server.continuouslyRegisterForwardedWorker(
							logger,
							channel,
							gardenForward.boundPort,
							0,
						)
						if err != nil {
							logger.Error("failed-to-register", err)
							return
						}

						processes = append(processes, process)

						break
					}

				case 2:
					gardenForward, found := forwards[r.gardenAddr]
					if !found {
						fmt.Fprintf(channel, "garden address %s not found in forwards\n", r.gardenAddr)
						return
					}

					baggageclaimForward, found := forwards[r.baggageclaimAddr]
					if !found {
						fmt.Fprintf(channel, "baggageclaim address %s not found in forwards\n", r.gardenAddr)
						return
					}

					process, err = server.continuouslyRegisterForwardedWorker(
						logger,
						channel,
						gardenForward.boundPort,
						baggageclaimForward.boundPort,
					)
					if err != nil {
						logger.Error("failed-to-register", err)
						return
					}

					processes = append(processes, process)
				}

				err = conn.Wait()
				logger.Error("connection-closed", err)

			default:
				logger.Info("invalid-command", lager.Data{
					"command": request.Command,
				})

				req.Reply(false, nil)
			}
		}
	}
}

func (server *registrarSSHServer) continuouslyRegisterWorkerDirectly(
	logger lager.Logger,
	channel ssh.Channel,
) (ifrit.Process, error) {
	logger.Session("start")
	defer logger.Session("done")

	var worker atc.Worker
	err := json.NewDecoder(channel).Decode(&worker)
	if err != nil {
		return nil, err
	}

	return server.heartbeatWorker(logger, worker, channel), nil
}

func (server *registrarSSHServer) continuouslyRegisterForwardedWorker(
	logger lager.Logger,
	channel ssh.Channel,
	gardenPort uint32,
	baggageclaimPort uint32,
) (ifrit.Process, error) {
	logger.Session("start")
	defer logger.Session("done")

	var worker atc.Worker
	err := json.NewDecoder(channel).Decode(&worker)
	if err != nil {
		return nil, err
	}

	worker.GardenAddr = fmt.Sprintf("%s:%d", server.forwardHost, gardenPort)

	if baggageclaimPort != 0 {
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", server.forwardHost, baggageclaimPort)
	}

	return server.heartbeatWorker(logger, worker, channel), nil
}

func (server *registrarSSHServer) heartbeatWorker(logger lager.Logger, worker atc.Worker, channel ssh.Channel) ifrit.Process {
	return ifrit.Background(tsa.NewHeartbeater(
		logger,
		clock.NewClock(),
		server.heartbeatInterval,
		server.cprInterval,
		gclient.New(gconn.NewWithDialerAndLogger(keepaliveDialerFactory("tcp", worker.GardenAddr), logger.Session("garden-connection"))),
		server.atcEndpoint,
		server.tokenGenerator,
		worker,
		channel,
	))
}

func (server *registrarSSHServer) handleForwardRequests(
	logger lager.Logger,
	conn *ssh.ServerConn,
	reqs <-chan *ssh.Request,
	forwardedTCPIPs chan<- forwardedTCPIP,
) {
	var forwardedThings int

	for r := range reqs {
		switch r.Type {
		case "tcpip-forward":
			logger := logger.Session("tcpip-forward")

			forwardedThings++

			if forwardedThings > 2 {
				logger.Info("rejecting-extra-forward-request")
				r.Reply(false, nil)
				continue
			}

			var req tcpipForwardRequest
			err := ssh.Unmarshal(r.Payload, &req)
			if err != nil {
				logger.Error("malformed-tcpip-request", err)
				r.Reply(false, nil)
				continue
			}

			listener, err := net.Listen("tcp", "0.0.0.0:0")
			if err != nil {
				logger.Error("failed-to-listen", err)
				r.Reply(false, nil)
				continue
			}

			defer listener.Close()

			bindAddr := net.JoinHostPort(req.BindIP, fmt.Sprintf("%d", req.BindPort))

			logger.Info("forwarding-tcpip", lager.Data{
				"requested-bind-addr": bindAddr,
			})

			_, port, err := net.SplitHostPort(listener.Addr().String())
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			var res tcpipForwardResponse
			_, err = fmt.Sscanf(port, "%d", &res.BoundPort)
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			forPort := req.BindPort
			if forPort == 0 {
				forPort = res.BoundPort
			}

			process := server.forwardTCPIP(logger, conn, listener, req.BindIP, forPort)

			forwardedTCPIPs <- forwardedTCPIP{
				bindAddr:  fmt.Sprintf("%s:%d", req.BindIP, req.BindPort),
				boundPort: res.BoundPort,
				process:   process,
			}

			r.Reply(true, ssh.Marshal(res))

		default:
			logger.Info("ignoring-request", lager.Data{"type": r.Type})
			r.Reply(false, nil)
		}
	}
}

func (server *registrarSSHServer) forwardTCPIP(
	logger lager.Logger,
	conn *ssh.ServerConn,
	listener net.Listener,
	forwardIP string,
	forwardPort uint32,
) ifrit.Process {
	return ifrit.Background(ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
		go func() {
			<-signals

			listener.Close()
		}()

		close(ready)

		for {
			localConn, err := listener.Accept()
			if err != nil {
				logger.Error("failed-to-accept", err)
				break
			}

			go forwardLocalConn(logger, localConn, conn, forwardIP, forwardPort)
		}

		return nil
	}))
}

func forwardLocalConn(logger lager.Logger, localConn net.Conn, conn *ssh.ServerConn, forwardIP string, forwardPort uint32) {
	defer localConn.Close()

	var req forwardTCPIPChannelRequest
	req.ForwardIP = forwardIP
	req.ForwardPort = forwardPort

	host, port, err := net.SplitHostPort(localConn.RemoteAddr().String())
	if err != nil {
		logger.Error("failed-to-split-host-port", err)
		return
	}

	req.OriginIP = host
	_, err = fmt.Sscanf(port, "%d", &req.OriginPort)
	if err != nil {
		logger.Error("failed-to-parse-port", err)
		return
	}

	channel, reqs, err := conn.OpenChannel("forwarded-tcpip", ssh.Marshal(req))
	if err != nil {
		logger.Error("failed-to-open-channel", err)
		return
	}

	defer channel.Close()

	go func() {
		for r := range reqs {
			logger.Info("ignoring-request", lager.Data{
				"type": r.Type,
			})

			r.Reply(false, nil)
		}
	}()

	wg := new(sync.WaitGroup)

	pipe := func(to io.WriteCloser, from io.ReadCloser) {
		// if either end breaks, close both ends to ensure they're both unblocked,
		// otherwise io.Copy can block forever if e.g. reading after write end has
		// gone away
		defer to.Close()
		defer from.Close()
		defer wg.Done()

		io.Copy(to, from)
	}

	wg.Add(1)
	go pipe(localConn, channel)

	wg.Add(1)
	go pipe(channel, localConn)

	wg.Wait()
}

func keepaliveDialerFactory(network string, address string) gconn.DialerFunc {
	return func(string, string) (net.Conn, error) {
		return keepaliveDialer(network, address)
	}
}
<commit_msg>reply successfully to keepalive requests<commit_after>package tsacmd

import (
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	gclient "github.com/cloudfoundry-incubator/garden/client"
	gconn "github.com/cloudfoundry-incubator/garden/client/connection"
	"github.com/concourse/atc"
	"github.com/concourse/tsa"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/rata"
	"golang.org/x/crypto/ssh"
)

type registrarSSHServer struct {
	logger            lager.Logger
	atcEndpoint       *rata.RequestGenerator
	tokenGenerator    tsa.TokenGenerator
	heartbeatInterval time.Duration
	cprInterval       time.Duration
	forwardHost       string
	config            *ssh.ServerConfig
	httpClient        *http.Client
}

func (server *registrarSSHServer) Serve(listener net.Listener) {
	for {
		c, err := listener.Accept()
		if err != nil {
			if !strings.Contains(err.Error(), "use of closed network connection") {
				server.logger.Error("failed-to-accept", err)
			}

			return
		}

		logger := server.logger.Session("connection")

		conn, chans, reqs, err := ssh.NewServerConn(c, server.config)
		if err != nil {
			logger.Info("handshake-failed", lager.Data{"error": err.Error()})
			continue
		}

		go server.handleConn(logger, conn, chans, reqs)
	}
}

type forwardedTCPIP struct {
	bindAddr  string
	process   ifrit.Process
	boundPort uint32
}

func (server *registrarSSHServer) handleConn(logger lager.Logger, conn *ssh.ServerConn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request) {
	defer conn.Close()

	forwardedTCPIPs := make(chan forwardedTCPIP, 2)
	go server.handleForwardRequests(logger, conn, reqs, forwardedTCPIPs)

	var processes []ifrit.Process
	var process ifrit.Process

	// ensure processes get cleaned up
	defer func() {
		cleanupLog := logger.Session("cleanup")

		for _, p := range processes {
			cleanupLog.Debug("interrupting")

			p.Signal(os.Interrupt)
		}

		for _, p := range processes {
			err := <-p.Wait()
			if err != nil {
				cleanupLog.Error("process-exited-with-failure", err)
			} else {
				cleanupLog.Debug("process-exited-successfully")
			}
		}
	}()

	for newChannel := range chans {
		if newChannel.ChannelType() != "session" {
			logger.Info("rejecting-unknown-channel-type", lager.Data{
				"type": newChannel.ChannelType(),
			})

			newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
			continue
		}

		channel, requests, err := newChannel.Accept()
		if err != nil {
			logger.Error("failed-to-accept-channel", err)
			return
		}

		defer channel.Close()

		for req := range requests {
			logger.Info("channel-request", lager.Data{
				"type": req.Type,
			})

			if req.Type != "exec" {
				logger.Info("rejecting")
				req.Reply(false, nil)
				continue
			}

			var request execRequest
			err = ssh.Unmarshal(req.Payload, &request)
			if err != nil {
				logger.Error("malformed-exec-request", err)
				req.Reply(false, nil)
				return
			}

			workerRequest, err := parseRequest(request.Command)
			if err != nil {
				fmt.Fprintf(channel, "invalid command: %s", err)
				req.Reply(false, nil)
				continue
			}

			switch r := workerRequest.(type) {
			case registerWorkerRequest:
				logger := logger.Session("register-worker")

				req.Reply(true, nil)

				process, err = server.continuouslyRegisterWorkerDirectly(logger, channel)
				if err != nil {
					logger.Error("failed-to-register", err)
					return
				}

				processes = append(processes, process)

				err = conn.Wait()
				logger.Error("connection-closed", err)

			case forwardWorkerRequest:
				logger := logger.Session("forward-worker")

				req.Reply(true, nil)

				forwards := map[string]forwardedTCPIP{}

				for i := 0; i < r.expectedForwards(); i++ {
					select {
					case forwarded := <-forwardedTCPIPs:
						logger.Info("forwarded-tcpip", lager.Data{
							"bound-port": forwarded.boundPort,
						})

						processes = append(processes, forwarded.process)

						forwards[forwarded.bindAddr] = forwarded

					case <-time.After(10 * time.Second): // todo better?
						logger.Info("never-forwarded-tcpip")
					}
				}

				switch len(forwards) {
				case 0:
					fmt.Fprintf(channel, "requested forwarding but no forwards given\n")
					return

				case 1:
					for _, gardenForward := range forwards {
						process, err = server.continuouslyRegisterForwardedWorker(
							logger,
							channel,
							gardenForward.boundPort,
							0,
						)
						if err != nil {
							logger.Error("failed-to-register", err)
							return
						}

						processes = append(processes, process)

						break
					}

				case 2:
					gardenForward, found := forwards[r.gardenAddr]
					if !found {
						fmt.Fprintf(channel, "garden address %s not found in forwards\n", r.gardenAddr)
						return
					}

					baggageclaimForward, found := forwards[r.baggageclaimAddr]
					if !found {
						fmt.Fprintf(channel, "baggageclaim address %s not found in forwards\n", r.gardenAddr)
						return
					}

					process, err = server.continuouslyRegisterForwardedWorker(
						logger,
						channel,
						gardenForward.boundPort,
						baggageclaimForward.boundPort,
					)
					if err != nil {
						logger.Error("failed-to-register", err)
						return
					}

					processes = append(processes, process)
				}

				err = conn.Wait()
				logger.Error("connection-closed", err)

			default:
				logger.Info("invalid-command", lager.Data{
					"command": request.Command,
				})

				req.Reply(false, nil)
			}
		}
	}
}

func (server *registrarSSHServer) continuouslyRegisterWorkerDirectly(
	logger lager.Logger,
	channel ssh.Channel,
) (ifrit.Process, error) {
	logger.Session("start")
	defer logger.Session("done")

	var worker atc.Worker
	err := json.NewDecoder(channel).Decode(&worker)
	if err != nil {
		return nil, err
	}

	return server.heartbeatWorker(logger, worker, channel), nil
}

func (server *registrarSSHServer) continuouslyRegisterForwardedWorker(
	logger lager.Logger,
	channel ssh.Channel,
	gardenPort uint32,
	baggageclaimPort uint32,
) (ifrit.Process, error) {
	logger.Session("start")
	defer logger.Session("done")

	var worker atc.Worker
	err := json.NewDecoder(channel).Decode(&worker)
	if err != nil {
		return nil, err
	}

	worker.GardenAddr = fmt.Sprintf("%s:%d", server.forwardHost, gardenPort)

	if baggageclaimPort != 0 {
		worker.BaggageclaimURL = fmt.Sprintf("http://%s:%d", server.forwardHost, baggageclaimPort)
	}

	return server.heartbeatWorker(logger, worker, channel), nil
}

func (server *registrarSSHServer) heartbeatWorker(logger lager.Logger, worker atc.Worker, channel ssh.Channel) ifrit.Process {
	return ifrit.Background(tsa.NewHeartbeater(
		logger,
		clock.NewClock(),
		server.heartbeatInterval,
		server.cprInterval,
		gclient.New(gconn.NewWithDialerAndLogger(keepaliveDialerFactory("tcp", worker.GardenAddr), logger.Session("garden-connection"))),
		server.atcEndpoint,
		server.tokenGenerator,
		worker,
		channel,
	))
}

func (server *registrarSSHServer) handleForwardRequests(
	logger lager.Logger,
	conn *ssh.ServerConn,
	reqs <-chan *ssh.Request,
	forwardedTCPIPs chan<- forwardedTCPIP,
) {
	var forwardedThings int

	for r := range reqs {
		switch r.Type {
		case "tcpip-forward":
			logger := logger.Session("tcpip-forward")

			forwardedThings++

			if forwardedThings > 2 {
				logger.Info("rejecting-extra-forward-request")
				r.Reply(false, nil)
				continue
			}

			var req tcpipForwardRequest
			err := ssh.Unmarshal(r.Payload, &req)
			if err != nil {
				logger.Error("malformed-tcpip-request", err)
				r.Reply(false, nil)
				continue
			}

			listener, err := net.Listen("tcp", "0.0.0.0:0")
			if err != nil {
				logger.Error("failed-to-listen", err)
				r.Reply(false, nil)
				continue
			}

			defer listener.Close()

			bindAddr := net.JoinHostPort(req.BindIP, fmt.Sprintf("%d", req.BindPort))

			logger.Info("forwarding-tcpip", lager.Data{
				"requested-bind-addr": bindAddr,
			})

			_, port, err := net.SplitHostPort(listener.Addr().String())
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			var res tcpipForwardResponse
			_, err = fmt.Sscanf(port, "%d", &res.BoundPort)
			if err != nil {
				r.Reply(false, nil)
				continue
			}

			forPort := req.BindPort
			if forPort == 0 {
				forPort = res.BoundPort
			}

			process := server.forwardTCPIP(logger, conn, listener, req.BindIP, forPort)

			forwardedTCPIPs <- forwardedTCPIP{
				bindAddr:  fmt.Sprintf("%s:%d", req.BindIP, req.BindPort),
				boundPort: res.BoundPort,
				process:   process,
			}

			r.Reply(true, ssh.Marshal(res))

		default:
			// OpenSSH sends keepalive@openssh.com, but there may be other clients;
			// just check for 'keepalive'
			if strings.Contains(r.Type, "keepalive") {
				logger.Info("keepalive", lager.Data{"type": r.Type})
				r.Reply(true, nil)
			} else {
				logger.Info("ignoring-request", lager.Data{"type": r.Type})
				r.Reply(false, nil)
			}
		}
	}
}

func (server *registrarSSHServer) forwardTCPIP(
	logger lager.Logger,
	conn *ssh.ServerConn,
	listener net.Listener,
	forwardIP string,
	forwardPort uint32,
) ifrit.Process {
	return ifrit.Background(ifrit.RunFunc(func(signals <-chan os.Signal, ready chan<- struct{}) error {
		go func() {
			<-signals

			listener.Close()
		}()

		close(ready)

		for {
			localConn, err := listener.Accept()
			if err != nil {
				logger.Error("failed-to-accept", err)
				break
			}

			go forwardLocalConn(logger, localConn, conn, forwardIP, forwardPort)
		}

		return nil
	}))
}

func forwardLocalConn(logger lager.Logger, localConn net.Conn, conn *ssh.ServerConn, forwardIP string, forwardPort uint32) {
	defer localConn.Close()

	var req forwardTCPIPChannelRequest
	req.ForwardIP = forwardIP
	req.ForwardPort = forwardPort

	host, port, err := net.SplitHostPort(localConn.RemoteAddr().String())
	if err != nil {
		logger.Error("failed-to-split-host-port", err)
		return
	}

	req.OriginIP = host
	_, err = fmt.Sscanf(port, "%d", &req.OriginPort)
	if err != nil {
		logger.Error("failed-to-parse-port", err)
		return
	}

	channel, reqs, err := conn.OpenChannel("forwarded-tcpip", ssh.Marshal(req))
	if err != nil {
		logger.Error("failed-to-open-channel", err)
		return
	}

	defer channel.Close()

	go func() {
		for r := range reqs {
			logger.Info("ignoring-request", lager.Data{
				"type": r.Type,
			})

			r.Reply(false, nil)
		}
	}()

	wg := new(sync.WaitGroup)

	pipe := func(to io.WriteCloser, from io.ReadCloser) {
		// if either end breaks, close both ends to ensure they're both unblocked,
		// otherwise io.Copy can block forever if e.g. reading after write end has
		// gone away
		defer to.Close()
		defer from.Close()
		defer wg.Done()

		io.Copy(to, from)
	}

	wg.Add(1)
	go pipe(localConn, channel)

	wg.Add(1)
	go pipe(channel, localConn)

	wg.Wait()
}

func keepaliveDialerFactory(network string, address string) gconn.DialerFunc {
	return func(string, string) (net.Conn, error) {
		return keepaliveDialer(network, address)
	}
}
<|endoftext|>


<commit_before>package qr

import (
	"fmt"
	"os"
	"strconv"
)

// A Version represents a QR version.
// The version specifies the size of the QR code:
// a QR code with version v has 4v+17 pixels on a side.
// Versions number from 1 to 40: the larger the version,
// the more information the code can store.
// A non-positive version means to select the
// version automatically.
type Version int

func (v Version) String() string {
	if v < 1 {
		return "auto"
	}
	return strconv.Itoa(int(v))
}

// A Mode represents a QR mode.
// The mode specifies the character set and, indirectly,
// the encoding.  The more precise the mode, the shorter
// the encoded data.
type Mode int

const (
	Numeric      Mode = 1
	Alphanumeric Mode = 2
	EightBit     Mode = 4
)

func (m Mode) String() string {
	switch m {
	case Numeric:
		return "numeric"
	case Alphanumeric:
		return "alpha"
	case EightBit:
		return "8bit"
	}
	return strconv.Itoa(int(m))
}

// A Pixel describes a single pixel in a QR code.
type Pixel uint32

const (
	Black Pixel = 1 << iota
	Invert
)

func (p Pixel) Offset() int {
	return int(p >> 5)
}

func OffsetPixel(o int) Pixel {
	return Pixel(o << 5)
}

func (r PixelRole) Pixel() Pixel {
	return Pixel(r << 2)
}

func (p Pixel) Role() PixelRole {
	return PixelRole(p>>2) & 7
}

func (p Pixel) String() string {
	s := p.Role().String()
	if p&Black != 0 {
		s += "+black"
	}
	if p&Invert != 0 {
		s += "+invert"
	}
	s += "+" + strconv.Itoa(p.Offset())
	return s
}

// A PixelRole describes the role of a QR pixel.
type PixelRole uint32

const (
	_         PixelRole = iota
	Position            // position squares (large)
	Alignment           // alignment squares (small)
	Timing              // timing strip between position squares
	Format              // format metadata
	Data                // data bit
	Check               // error correction check bit
)

var roles = []string{
	"",
	"position",
	"alignment",
	"timing",
	"format",
	"data",
	"check",
}

func (r PixelRole) String() string {
	if Position <= r && r <= Check {
		return roles[r]
	}
	return strconv.Itoa(int(r))
}

// A Level represents a QR error correction level.
// From least to most tolerant of errors, they are L, M, Q, H.
type Level int

const (
	L Level = 0
	M
	Q
	H
)

func (l Level) String() string {
	if L <= l && l <= H {
		return "LMQH"[l : l+1]
	}
	return strconv.Itoa(int(l))
}

// A Mask describes a mask that is applied to the QR
// code to avoid QR artifacts being interpreted as
// alignment and timing patterns (such as the squares
// in the corners).
type Mask int

// TODO: fill in masks

// A Plan describes how to construct a QR code
// with a specific version, level, and mask.
type Plan struct {
	Version Version
	Level   Level
	Mask    Mask

	DataBytes  int // number of data bytes
	CheckBytes int // number of error correcting (checksum) bytes
	Blocks     int // number of data blocks

	Pixel [][]Pixel // pixel map
}

// NewPlan returns a Plan for a QR code with the given
// version, level, and mask.
func NewPlan(version Version, level Level, mask Mask) (*Plan, os.Error) {
	p, err := vplan(version)
	if err != nil {
		return nil, err
	}
	if err := fplan(level, mask, p); err != nil {
		return nil, err
	}
	if err := lplan(level, p); err != nil {
		return nil, err
	}
	if err := mplan(mask, p); err != nil {
		return nil, err
	}
	return p, nil
}

// A version describes metadata associated with a version.
type version struct {
	apos    int
	astride int
}

var vtab = []version{
	{}, // dummy version 0
	// 1-10
	{100, 100},
	{16, 100},
	{20, 100},
	{24, 100},
	{28, 100},
	{32, 100},
	{20, 16},
	{22, 18},
	{24, 20},
	{26, 22},
	// 11-20
	{28, 24},
	{30, 26},
	{32, 28},
	{24, 20},
	{24, 22},
	{24, 24},
	{28, 24},
	{28, 26},
	{28, 28},
	{32, 28},
	// 21-30
	{26, 22},
	{24, 24},
	{28, 24},
	{26, 26},
	{30, 26},
	{28, 28},
	{32, 28},
	{24, 24},
	{28, 24},
	{24, 26},
	// 31-40
	{28, 26},
	{32, 26},
	{28, 28},
	{32, 28},
	{28, 24},
	{22, 26},
	{26, 26},
	{30, 26},
	{24, 28},
	{28, 28},
}

// vplan creates a Plan for the given version.
func vplan(v Version) (*Plan, os.Error) {
	p := &Plan{Version: v}
	if v < 1 || v > 40 {
		return nil, fmt.Errorf("invalid QR version %d", int(v))
	}
	siz := 17 + int(v)*4
	m := make([][]Pixel, siz)
	pix := make([]Pixel, siz*siz)
	for i := range m {
		m[i], pix = pix[:siz], pix[siz:]
	}
	p.Pixel = m

	// Timing markers (overwritten by boxes).
	// TODO: are there more in higher versions?
	const ti = 6 // timing is in row/column 6 (counting from 0)
	for i := range m {
		p := Timing.Pixel()
		if i&1 == 0 {
			p |= Black
		}
		m[i][ti] = p
		m[ti][i] = p
	}

	// Position boxes.
	posBox(m, 0, 0)
	posBox(m, siz-7, 0)
	posBox(m, 0, siz-7)

	// Alignment boxes.
	info := &vtab[v]
	for x := 4; x+5 < siz; {
		for y := 4; y+5 < siz; {
			// don't overwrite timing markers
			if (x < 7 && y < 7) || (x < 7 && y+5 >= siz-7) || (x+5 >= siz-7 && y < 7) {
			} else {
				alignBox(m, x, y)
			}
			if y == 4 {
				y = info.apos
			} else {
				y += info.astride
			}
		}
		if x == 4 {
			x = info.apos
		} else {
			x += info.astride
		}
	}
	return p, nil
}

// fplan adds the format pixels
func fplan(l Level, m Mask, p *Plan) os.Error {
	// Format pixels.
	fb := uint32(l^1) << 13 // level: L=01, M=00, Q=11, H=10
	fb |= uint32(m) << 10   // mask
	const formatPoly = 0x537
	rem := fb
	for i := 14; i >= 10; i-- {
		if rem&(1<<uint(i)) != 0 {
			rem ^= formatPoly << uint(i-10)
		}
	}
	fb |= rem
	invert := uint32(0x5412)
	siz := len(p.Pixel)
	for i := 0; i < 15; i++ {
		pix := Format.Pixel() + OffsetPixel(i)
		if (fb>>uint(i))&1 == 1 {
			pix |= Black
		}
		if (invert>>uint(i))&1 == 1 {
			pix ^= Invert | Black
		}
		// top left
		switch {
		case i < 6:
			p.Pixel[i][8] = pix
		case i < 8:
			p.Pixel[i+1][8] = pix
		case i < 9:
			p.Pixel[8][7] = pix
		default:
			p.Pixel[8][14-i] = pix
		}
		// bottom right
		switch {
		case i < 8:
			p.Pixel[8][siz-1-i] = pix
		default:
			p.Pixel[siz-1-(14-i)][8] = pix
		}
	}
	return nil
}

// lplan edits a version-only Plan to add information
// about the error correction levels.
func lplan(l Level, p *Plan) os.Error {
	p.Level = l

	// TODO: prepare actual pix
	siz := len(p.Pixel)
	pix := make([]Pixel, siz*siz)

	// Sweep up pair of columns,
	// then down, assigning to right then left pixel.
	// Repeat.
	// See Figure 2 of http://www.pclviewer.com/rs2/qrtopology.htm
	for x := siz; x > 0; {
		for y := siz-1; y >= 0; y-- {
			if p.Pixel[y][x-1].Role() == 0 {
				p.Pixel[y][x-1], pix = pix[0], pix[1:]
			}
			if p.Pixel[y][x-2].Role() == 0 {
				p.Pixel[y][x-2], pix = pix[0], pix[1:]
			}
		}
		x -= 2
		if x == 7 {  // vertical timing strip
			x--
		}
		for y := 0; y < siz; y++ {
			if p.Pixel[y][x-1].Role() == 0 {
				p.Pixel[y][x-1], pix = pix[0], pix[1:]
			}
			if p.Pixel[y][x-2].Role() == 0 {
				p.Pixel[y][x-2], pix = pix[0], pix[1:]
			}
		}
		x -= 2
	}

	return nil
}

// http://www.swetake.com/qr/qr5_en.html
var mfunc = []func(int, int) bool{
	func(i, j int) bool { return (i+j)%2 == 0 },
	func(i, j int) bool { return i%2 == 0 },
	func(i, j int) bool { return j%3 == 0 },
	func(i, j int) bool { return (i+j)%3 == 0 },
	func(i, j int) bool { return (i/2+j/3)%2 == 0 },
	func(i, j int) bool { return i*j%2+i*j%3 == 0 },
	func(i, j int) bool { return (i*j%2+i*j%3)%2 == 0 },
	func(i, j int) bool { return (i*j%3+(i+j)%2)%2 == 0 },
}

// mplan edits a version+level-only Plan to add the mask.
func mplan(m Mask, p *Plan) os.Error {
	f := mfunc[m]
	p.Mask = m
	for y, row := range p.Pixel {
		for x, pix := range row {
			if r := pix.Role(); (r == Data || r == Check) && f(x, y) {
				row[x] ^= Black|Invert
			}
		}
	}
	return nil
}

// posBox draws a position (large) box at upper left x, y.
func posBox(m [][]Pixel, x, y int) {
	pos := Position.Pixel()
	// box
	for dy := 0; dy < 7; dy++ {
		for dx := 0; dx < 7; dx++ {
			p := pos
			if dx == 0 || dx == 6 || dy == 0 || dy == 6 || 2 <= dx && dx <= 4 && 2 <= dy && dy <= 4 {
				p |= Black
			}
			m[y+dy][x+dx] = p
		}
	}
	// white border
	for dy := -1; dy < 8; dy++ {
		if 0 <= y+dy && y+dy < len(m) {
			if x > 0 {
				m[y+dy][x-1] = pos
			}
			if x+7 < len(m) {
				m[y+dy][x+7] = pos
			}
		}
	}
	for dx := -1; dx < 8; dx++ {
		if 0 <= x+dx && x+dx < len(m) {
			if y > 0 {
				m[y-1][x+dx] = pos
			}
			if y+7 < len(m) {
				m[y+7][x+dx] = pos
			}
		}
	}
}

// alignBox draw an alignment (small) box at upper left x, y.
func alignBox(m [][]Pixel, x, y int) {
	// box
	align := Alignment.Pixel()
	for dy := 0; dy < 5; dy++ {
		for dx := 0; dx < 5; dx++ {
			p := align
			if dx == 0 || dx == 4 || dy == 0 || dy == 4 || dx == 2 && dy == 2 {
				p |= Black
			}
			m[y+dy][x+dx] = p
		}
	}
}
<commit_msg>gofmt<commit_after>package qr

import (
	"fmt"
	"os"
	"strconv"
)

// A Version represents a QR version.
// The version specifies the size of the QR code:
// a QR code with version v has 4v+17 pixels on a side.
// Versions number from 1 to 40: the larger the version,
// the more information the code can store.
// A non-positive version means to select the
// version automatically.
type Version int

func (v Version) String() string {
	if v < 1 {
		return "auto"
	}
	return strconv.Itoa(int(v))
}

// A Mode represents a QR mode.
// The mode specifies the character set and, indirectly,
// the encoding.  The more precise the mode, the shorter
// the encoded data.
type Mode int

const (
	Numeric      Mode = 1
	Alphanumeric Mode = 2
	EightBit     Mode = 4
)

func (m Mode) String() string {
	switch m {
	case Numeric:
		return "numeric"
	case Alphanumeric:
		return "alpha"
	case EightBit:
		return "8bit"
	}
	return strconv.Itoa(int(m))
}

// A Pixel describes a single pixel in a QR code.
type Pixel uint32

const (
	Black Pixel = 1 << iota
	Invert
)

func (p Pixel) Offset() int {
	return int(p >> 5)
}

func OffsetPixel(o int) Pixel {
	return Pixel(o << 5)
}

func (r PixelRole) Pixel() Pixel {
	return Pixel(r << 2)
}

func (p Pixel) Role() PixelRole {
	return PixelRole(p>>2) & 7
}

func (p Pixel) String() string {
	s := p.Role().String()
	if p&Black != 0 {
		s += "+black"
	}
	if p&Invert != 0 {
		s += "+invert"
	}
	s += "+" + strconv.Itoa(p.Offset())
	return s
}

// A PixelRole describes the role of a QR pixel.
type PixelRole uint32

const (
	_         PixelRole = iota
	Position            // position squares (large)
	Alignment           // alignment squares (small)
	Timing              // timing strip between position squares
	Format              // format metadata
	Data                // data bit
	Check               // error correction check bit
)

var roles = []string{
	"",
	"position",
	"alignment",
	"timing",
	"format",
	"data",
	"check",
}

func (r PixelRole) String() string {
	if Position <= r && r <= Check {
		return roles[r]
	}
	return strconv.Itoa(int(r))
}

// A Level represents a QR error correction level.
// From least to most tolerant of errors, they are L, M, Q, H.
type Level int

const (
	L Level = 0
	M
	Q
	H
)

func (l Level) String() string {
	if L <= l && l <= H {
		return "LMQH"[l : l+1]
	}
	return strconv.Itoa(int(l))
}

// A Mask describes a mask that is applied to the QR
// code to avoid QR artifacts being interpreted as
// alignment and timing patterns (such as the squares
// in the corners).
type Mask int

// TODO: fill in masks

// A Plan describes how to construct a QR code
// with a specific version, level, and mask.
type Plan struct {
	Version Version
	Level   Level
	Mask    Mask

	DataBytes  int // number of data bytes
	CheckBytes int // number of error correcting (checksum) bytes
	Blocks     int // number of data blocks

	Pixel [][]Pixel // pixel map
}

// NewPlan returns a Plan for a QR code with the given
// version, level, and mask.
func NewPlan(version Version, level Level, mask Mask) (*Plan, os.Error) {
	p, err := vplan(version)
	if err != nil {
		return nil, err
	}
	if err := fplan(level, mask, p); err != nil {
		return nil, err
	}
	if err := lplan(level, p); err != nil {
		return nil, err
	}
	if err := mplan(mask, p); err != nil {
		return nil, err
	}
	return p, nil
}

// A version describes metadata associated with a version.
type version struct {
	apos    int
	astride int
}

var vtab = []version{
	{}, // dummy version 0
	// 1-10
	{100, 100},
	{16, 100},
	{20, 100},
	{24, 100},
	{28, 100},
	{32, 100},
	{20, 16},
	{22, 18},
	{24, 20},
	{26, 22},
	// 11-20
	{28, 24},
	{30, 26},
	{32, 28},
	{24, 20},
	{24, 22},
	{24, 24},
	{28, 24},
	{28, 26},
	{28, 28},
	{32, 28},
	// 21-30
	{26, 22},
	{24, 24},
	{28, 24},
	{26, 26},
	{30, 26},
	{28, 28},
	{32, 28},
	{24, 24},
	{28, 24},
	{24, 26},
	// 31-40
	{28, 26},
	{32, 26},
	{28, 28},
	{32, 28},
	{28, 24},
	{22, 26},
	{26, 26},
	{30, 26},
	{24, 28},
	{28, 28},
}

// vplan creates a Plan for the given version.
func vplan(v Version) (*Plan, os.Error) {
	p := &Plan{Version: v}
	if v < 1 || v > 40 {
		return nil, fmt.Errorf("invalid QR version %d", int(v))
	}
	siz := 17 + int(v)*4
	m := make([][]Pixel, siz)
	pix := make([]Pixel, siz*siz)
	for i := range m {
		m[i], pix = pix[:siz], pix[siz:]
	}
	p.Pixel = m

	// Timing markers (overwritten by boxes).
	// TODO: are there more in higher versions?
	const ti = 6 // timing is in row/column 6 (counting from 0)
	for i := range m {
		p := Timing.Pixel()
		if i&1 == 0 {
			p |= Black
		}
		m[i][ti] = p
		m[ti][i] = p
	}

	// Position boxes.
	posBox(m, 0, 0)
	posBox(m, siz-7, 0)
	posBox(m, 0, siz-7)

	// Alignment boxes.
	info := &vtab[v]
	for x := 4; x+5 < siz; {
		for y := 4; y+5 < siz; {
			// don't overwrite timing markers
			if (x < 7 && y < 7) || (x < 7 && y+5 >= siz-7) || (x+5 >= siz-7 && y < 7) {
			} else {
				alignBox(m, x, y)
			}
			if y == 4 {
				y = info.apos
			} else {
				y += info.astride
			}
		}
		if x == 4 {
			x = info.apos
		} else {
			x += info.astride
		}
	}
	return p, nil
}

// fplan adds the format pixels
func fplan(l Level, m Mask, p *Plan) os.Error {
	// Format pixels.
	fb := uint32(l^1) << 13 // level: L=01, M=00, Q=11, H=10
	fb |= uint32(m) << 10   // mask
	const formatPoly = 0x537
	rem := fb
	for i := 14; i >= 10; i-- {
		if rem&(1<<uint(i)) != 0 {
			rem ^= formatPoly << uint(i-10)
		}
	}
	fb |= rem
	invert := uint32(0x5412)
	siz := len(p.Pixel)
	for i := 0; i < 15; i++ {
		pix := Format.Pixel() + OffsetPixel(i)
		if (fb>>uint(i))&1 == 1 {
			pix |= Black
		}
		if (invert>>uint(i))&1 == 1 {
			pix ^= Invert | Black
		}
		// top left
		switch {
		case i < 6:
			p.Pixel[i][8] = pix
		case i < 8:
			p.Pixel[i+1][8] = pix
		case i < 9:
			p.Pixel[8][7] = pix
		default:
			p.Pixel[8][14-i] = pix
		}
		// bottom right
		switch {
		case i < 8:
			p.Pixel[8][siz-1-i] = pix
		default:
			p.Pixel[siz-1-(14-i)][8] = pix
		}
	}
	return nil
}

// lplan edits a version-only Plan to add information
// about the error correction levels.
func lplan(l Level, p *Plan) os.Error {
	p.Level = l

	// TODO: prepare actual pix
	siz := len(p.Pixel)
	pix := make([]Pixel, siz*siz)

	// Sweep up pair of columns,
	// then down, assigning to right then left pixel.
	// Repeat.
	// See Figure 2 of http://www.pclviewer.com/rs2/qrtopology.htm
	for x := siz; x > 0; {
		for y := siz - 1; y >= 0; y-- {
			if p.Pixel[y][x-1].Role() == 0 {
				p.Pixel[y][x-1], pix = pix[0], pix[1:]
			}
			if p.Pixel[y][x-2].Role() == 0 {
				p.Pixel[y][x-2], pix = pix[0], pix[1:]
			}
		}
		x -= 2
		if x == 7 { // vertical timing strip
			x--
		}
		for y := 0; y < siz; y++ {
			if p.Pixel[y][x-1].Role() == 0 {
				p.Pixel[y][x-1], pix = pix[0], pix[1:]
			}
			if p.Pixel[y][x-2].Role() == 0 {
				p.Pixel[y][x-2], pix = pix[0], pix[1:]
			}
		}
		x -= 2
	}

	return nil
}

// http://www.swetake.com/qr/qr5_en.html
var mfunc = []func(int, int) bool{
	func(i, j int) bool { return (i+j)%2 == 0 },
	func(i, j int) bool { return i%2 == 0 },
	func(i, j int) bool { return j%3 == 0 },
	func(i, j int) bool { return (i+j)%3 == 0 },
	func(i, j int) bool { return (i/2+j/3)%2 == 0 },
	func(i, j int) bool { return i*j%2+i*j%3 == 0 },
	func(i, j int) bool { return (i*j%2+i*j%3)%2 == 0 },
	func(i, j int) bool { return (i*j%3+(i+j)%2)%2 == 0 },
}

// mplan edits a version+level-only Plan to add the mask.
func mplan(m Mask, p *Plan) os.Error {
	f := mfunc[m]
	p.Mask = m
	for y, row := range p.Pixel {
		for x, pix := range row {
			if r := pix.Role(); (r == Data || r == Check) && f(x, y) {
				row[x] ^= Black | Invert
			}
		}
	}
	return nil
}

// posBox draws a position (large) box at upper left x, y.
func posBox(m [][]Pixel, x, y int) {
	pos := Position.Pixel()
	// box
	for dy := 0; dy < 7; dy++ {
		for dx := 0; dx < 7; dx++ {
			p := pos
			if dx == 0 || dx == 6 || dy == 0 || dy == 6 || 2 <= dx && dx <= 4 && 2 <= dy && dy <= 4 {
				p |= Black
			}
			m[y+dy][x+dx] = p
		}
	}
	// white border
	for dy := -1; dy < 8; dy++ {
		if 0 <= y+dy && y+dy < len(m) {
			if x > 0 {
				m[y+dy][x-1] = pos
			}
			if x+7 < len(m) {
				m[y+dy][x+7] = pos
			}
		}
	}
	for dx := -1; dx < 8; dx++ {
		if 0 <= x+dx && x+dx < len(m) {
			if y > 0 {
				m[y-1][x+dx] = pos
			}
			if y+7 < len(m) {
				m[y+7][x+dx] = pos
			}
		}
	}
}

// alignBox draw an alignment (small) box at upper left x, y.
func alignBox(m [][]Pixel, x, y int) {
	// box
	align := Alignment.Pixel()
	for dy := 0; dy < 5; dy++ {
		for dx := 0; dx < 5; dx++ {
			p := align
			if dx == 0 || dx == 4 || dy == 0 || dy == 4 || dx == 2 && dy == 2 {
				p |= Black
			}
			m[y+dy][x+dx] = p
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"math"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/dustin/go-couchstore"
	"github.com/dustin/go-jsonpointer"
)

var timeoutError = errors.New("query timed out")

type ptrval struct {
	di       *couchstore.DocInfo
	val      *string
	included bool
}

type Reducer func(input chan ptrval) interface{}

type processOut struct {
	cacheKey    string
	key         int64
	value       []interface{}
	err         error
	cacheOpaque uint32
}

func (p processOut) MarshalJSON() ([]byte, error) {
	return json.Marshal(map[string]interface{}{"v": p.value})
}

type processIn struct {
	cacheKey   string
	dbname     string
	key        int64
	infos      []*couchstore.DocInfo
	nextInfo   *couchstore.DocInfo
	ptrs       []string
	reds       []string
	before     time.Time
	filters    []string
	filtervals []string
	out        chan<- *processOut
}

type queryIn struct {
	dbname     string
	from       string
	to         string
	group      int
	ptrs       []string
	reds       []string
	start      time.Time
	before     time.Time
	filters    []string
	filtervals []string
	started    int32
	totalKeys  int32
	out        chan *processOut
	cherr      chan error
}

func processDoc(di *couchstore.DocInfo, chs []chan ptrval,
	doc []byte, ptrs []string,
	filters []string, filtervals []string,
	included bool) {

	pv := ptrval{di, nil, included}

	j := map[string]interface{}{}
	err := json.Unmarshal(doc, &j)
	if err != nil {
		for i := range ptrs {
			chs[i] <- pv
		}
		return
	}
	for i, p := range filters {
		val := jsonpointer.Get(j, p)
		checkVal := filtervals[i]
		switch val.(type) {
		case string:
			if val != checkVal {
				return
			}
		case int, uint, int64, float64, uint64, bool:
			v := fmt.Sprintf("%v", val)
			if v != checkVal {
				return
			}
		default:
			return
		}
	}

	for i, p := range ptrs {
		val := jsonpointer.Get(j, p)
		switch x := val.(type) {
		case string:
			pv.val = &x
			chs[i] <- pv
		case int, uint, int64, float64, uint64, bool:
			v := fmt.Sprintf("%v", val)
			pv.val = &v
			chs[i] <- pv
		default:
			log.Printf("Ignoring %T", val)
			chs[i] <- pv
		}
	}
}

func process_docs(pi *processIn) {

	result := processOut{pi.cacheKey, pi.key, nil, nil, 0}

	if len(pi.ptrs) == 0 {
		log.Panicf("No pointers specified in query: %#v", *pi)
	}

	db, err := dbopen(pi.dbname)
	if err != nil {
		result.err = err
		pi.out <- &result
		return
	}
	defer db.Close()

	chans := make([]chan ptrval, 0, len(pi.ptrs))
	resultchs := make([]chan interface{}, 0, len(pi.ptrs))
	for i, r := range pi.reds {
		chans = append(chans, make(chan ptrval))
		resultchs = append(resultchs, make(chan interface{}))

		go func() {
			resultchs[i] <- reducers[r](chans[i])
		}()
	}

	go func() {
		defer closeAll(chans)

		dodoc := func(di *couchstore.DocInfo, included bool) {
			doc, err := db.GetFromDocInfo(di)
			if err == nil {
				processDoc(di, chans, doc.Value(), pi.ptrs,
					pi.filters, pi.filtervals, included)
			} else {
				for i := range pi.ptrs {
					chans[i] <- ptrval{di, nil, included}
				}
			}
		}

		for _, di := range pi.infos {
			dodoc(di, true)
		}
		if pi.nextInfo != nil {
			dodoc(pi.nextInfo, false)
		}
	}()

	results := make([]interface{}, len(pi.ptrs))
	for i := 0; i < len(pi.ptrs); i++ {
		results[i] = <-resultchs[i]
		if f, fok := results[i].(float64); fok && math.IsNaN(f) {
			results[i] = nil
		}
	}
	result.value = results

	if result.cacheOpaque == 0 && result.cacheKey != "" {
		// It's OK if we can't store our newly pulled item in
		// the cache, but it's most definitely not OK to stop
		// here because of this.
		select {
		case cacheInputSet <- &result:
		default:
		}
	}
	pi.out <- &result
}

func docProcessor(ch <-chan *processIn) {
	for pi := range ch {
		if time.Now().Before(pi.before) {
			process_docs(pi)
		} else {
			pi.out <- &processOut{"", pi.key, nil, timeoutError, 0}
		}
	}
}

func fetchDocs(dbname string, key int64, infos []*couchstore.DocInfo,
	nextInfo *couchstore.DocInfo, ptrs []string, reds []string,
	filters []string, filtervals []string,
	before time.Time, out chan<- *processOut) {

	i := processIn{"", dbname, key, infos, nextInfo,
		ptrs, reds, before, filters, filtervals, out}

	cacheInput <- &i
}

func runQuery(q *queryIn) {
	if len(q.ptrs) == 0 {
		q.cherr <- fmt.Errorf("At least one pointer is required")
		return
	}

	db, err := dbopen(q.dbname)
	if err != nil {
		log.Printf("Error opening db: %v - %v", q.dbname, err)
		q.cherr <- err
		return
	}
	defer db.Close()

	chunk := int64(time.Duration(q.group) * time.Millisecond)

	infos := []*couchstore.DocInfo{}
	g := int64(0)
	nextg := ""

	err = db.Walk(q.from, func(d *couchstore.Couchstore,
		di *couchstore.DocInfo) error {
		kstr := di.ID()
		var err error
		if q.to != "" && kstr >= q.to {
			err = couchstore.StopIteration
		}

		atomic.AddInt32(&q.totalKeys, 1)

		if kstr >= nextg {
			if len(infos) > 0 {
				atomic.AddInt32(&q.started, 1)
				fetchDocs(q.dbname, g, infos, di,
					q.ptrs, q.reds, q.filters, q.filtervals,
					q.before, q.out)

				infos = make([]*couchstore.DocInfo, 0, len(infos))
			}

			k := parseKey(kstr)
			g = (k / chunk) * chunk
			nextgi := g + chunk
			nextgt := time.Unix(nextgi/1e9, nextgi%1e9).UTC()
			nextg = nextgt.Format(time.RFC3339Nano)
		}
		infos = append(infos, di)

		return err
	})

	if err == nil && len(infos) > 0 {
		atomic.AddInt32(&q.started, 1)
		fetchDocs(q.dbname, g, infos, nil,
			q.ptrs, q.reds, q.filters, q.filtervals,
			q.before, q.out)
	}

	q.cherr <- err
}

func queryExecutor(ch <-chan *queryIn) {
	for q := range ch {
		if time.Now().Before(q.before) {
			runQuery(q)
		} else {
			log.Printf("Timed out query that's %v late",
				time.Since(q.before))
			q.cherr <- timeoutError
		}
	}
}

func executeQuery(dbname, from, to string, group int,
	ptrs []string, reds []string, filters []string, filtervals []string) *queryIn {

	now := time.Now()

	rv := &queryIn{
		dbname:     dbname,
		from:       from,
		to:         to,
		group:      group,
		ptrs:       ptrs,
		reds:       reds,
		start:      now,
		before:     now.Add(*queryTimeout),
		filters:    filters,
		filtervals: filtervals,
		out:        make(chan *processOut),
		cherr:      make(chan error),
	}
	queryInput <- rv
	return rv
}

var processorInput chan *processIn
var queryInput chan *queryIn

func convertTofloat64(in chan ptrval) chan float64 {
	ch := make(chan float64)
	go func() {
		defer close(ch)
		for v := range in {
			if v.included && v.val != nil {
				x, err := strconv.ParseFloat(*v.val, 64)
				if err == nil {
					ch <- x
				}
			}
		}
	}()

	return ch
}

func convertTofloat64Rate(in chan ptrval) chan float64 {
	ch := make(chan float64)
	go func() {
		defer close(ch)
		var prevts int64
		var preval float64

		// First, find a part of the stream that has usable data.
		for v := range in {
			if v.di != nil && v.val != nil {
				x, err := strconv.ParseFloat(*v.val, 64)
				if err == nil {
					prevts = parseKey(v.di.ID())
					preval = x
					break
				}
			}
		}
		// Then emit floats based on deltas from previous values.
		for v := range in {
			if v.di != nil && v.val != nil {
				x, err := strconv.ParseFloat(*v.val, 64)
				if err == nil {
					thists := parseKey(v.di.ID())

					val := ((x - preval) /
						(float64(thists-prevts) / 1e9))

					if !math.IsNaN(val) {
						ch <- val
					}

					prevts = thists
					preval = x
				}
			}
		}
	}()

	return ch
}

var reducers = map[string]Reducer{
	"identity": func(input chan ptrval) interface{} {
		rv := []*string{}
		for s := range input {
			if s.included {
				rv = append(rv, s.val)
			}
		}
		return rv
	},
	"any": func(input chan ptrval) interface{} {
		for v := range input {
			if v.included && v.val != nil {
				return *v.val
			}
		}
		return nil
	},
	"count": func(input chan ptrval) interface{} {
		rv := 0
		for v := range input {
			if v.included && v.val != nil {
				rv++
			}
		}
		return rv
	},
	"sum": func(input chan ptrval) interface{} {
		rv := float64(0)
		for v := range convertTofloat64(input) {
			rv += v
		}
		return rv
	},
	"sumsq": func(input chan ptrval) interface{} {
		rv := float64(0)
		for v := range convertTofloat64(input) {
			rv += (v * v)
		}
		return rv
	},
	"max": func(input chan ptrval) interface{} {
		rv := math.NaN()
		for v := range convertTofloat64(input) {
			if v > rv || math.IsNaN(rv) {
				rv = v
			}
		}
		return rv
	},
	"min": func(input chan ptrval) interface{} {
		rv := math.NaN()
		for v := range convertTofloat64(input) {
			if v < rv || math.IsNaN(rv) {
				rv = v
			}
		}
		return rv
	},
	"avg": func(input chan ptrval) interface{} {
		nums := float64(0)
		sum := float64(0)
		for v := range convertTofloat64(input) {
			nums++
			sum += v
		}
		if nums > 0 {
			return sum / nums
		}
		return math.NaN()
	},
	"c_min": func(input chan ptrval) interface{} {
		rv := math.NaN()
		for v := range convertTofloat64Rate(input) {
			if v < rv || math.IsNaN(rv) {
				rv = v
			}
		}
		return rv
	},
	"c_avg": func(input chan ptrval) interface{} {
		nums := float64(0)
		sum := float64(0)
		for v := range convertTofloat64Rate(input) {
			nums++
			sum += v
		}
		if nums > 0 {
			return sum / nums
		}
		return math.NaN()
	},
	"c_max": func(input chan ptrval) interface{} {
		rv := math.NaN()
		for v := range convertTofloat64Rate(input) {
			if v > rv || math.IsNaN(rv) {
				rv = v
			}
		}
		return rv
	},
}
<commit_msg>Closing on loop parameters doesn't automatically copy them.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"math"
	"strconv"
	"sync/atomic"
	"time"

	"github.com/dustin/go-couchstore"
	"github.com/dustin/go-jsonpointer"
)

var timeoutError = errors.New("query timed out")

type ptrval struct {
	di       *couchstore.DocInfo
	val      *string
	included bool
}

type Reducer func(input chan ptrval) interface{}

type processOut struct {
	cacheKey    string
	key         int64
	value       []interface{}
	err         error
	cacheOpaque uint32
}

func (p processOut) MarshalJSON() ([]byte, error) {
	return json.Marshal(map[string]interface{}{"v": p.value})
}

type processIn struct {
	cacheKey   string
	dbname     string
	key        int64
	infos      []*couchstore.DocInfo
	nextInfo   *couchstore.DocInfo
	ptrs       []string
	reds       []string
	before     time.Time
	filters    []string
	filtervals []string
	out        chan<- *processOut
}

type queryIn struct {
	dbname     string
	from       string
	to         string
	group      int
	ptrs       []string
	reds       []string
	start      time.Time
	before     time.Time
	filters    []string
	filtervals []string
	started    int32
	totalKeys  int32
	out        chan *processOut
	cherr      chan error
}

func processDoc(di *couchstore.DocInfo, chs []chan ptrval,
	doc []byte, ptrs []string,
	filters []string, filtervals []string,
	included bool) {

	pv := ptrval{di, nil, included}

	j := map[string]interface{}{}
	err := json.Unmarshal(doc, &j)
	if err != nil {
		for i := range ptrs {
			chs[i] <- pv
		}
		return
	}
	for i, p := range filters {
		val := jsonpointer.Get(j, p)
		checkVal := filtervals[i]
		switch val.(type) {
		case string:
			if val != checkVal {
				return
			}
		case int, uint, int64, float64, uint64, bool:
			v := fmt.Sprintf("%v", val)
			if v != checkVal {
				return
			}
		default:
			return
		}
	}

	for i, p := range ptrs {
		val := jsonpointer.Get(j, p)
		switch x := val.(type) {
		case string:
			pv.val = &x
			chs[i] <- pv
		case int, uint, int64, float64, uint64, bool:
			v := fmt.Sprintf("%v", val)
			pv.val = &v
			chs[i] <- pv
		default:
			log.Printf("Ignoring %T", val)
			chs[i] <- pv
		}
	}
}

func process_docs(pi *processIn) {

	result := processOut{pi.cacheKey, pi.key, nil, nil, 0}

	if len(pi.ptrs) == 0 {
		log.Panicf("No pointers specified in query: %#v", *pi)
	}

	db, err := dbopen(pi.dbname)
	if err != nil {
		result.err = err
		pi.out <- &result
		return
	}
	defer db.Close()

	chans := make([]chan ptrval, 0, len(pi.ptrs))
	resultchs := make([]chan interface{}, 0, len(pi.ptrs))
	for i, r := range pi.reds {
		chans = append(chans, make(chan ptrval))
		resultchs = append(resultchs, make(chan interface{}))

		go func(fi int, fr string) {
			resultchs[fi] <- reducers[fr](chans[fi])
		}(i, r)
	}

	go func() {
		defer closeAll(chans)

		dodoc := func(di *couchstore.DocInfo, included bool) {
			doc, err := db.GetFromDocInfo(di)
			if err == nil {
				processDoc(di, chans, doc.Value(), pi.ptrs,
					pi.filters, pi.filtervals, included)
			} else {
				for i := range pi.ptrs {
					chans[i] <- ptrval{di, nil, included}
				}
			}
		}

		for _, di := range pi.infos {
			dodoc(di, true)
		}
		if pi.nextInfo != nil {
			dodoc(pi.nextInfo, false)
		}
	}()

	results := make([]interface{}, len(pi.ptrs))
	for i := 0; i < len(pi.ptrs); i++ {
		results[i] = <-resultchs[i]
		if f, fok := results[i].(float64); fok && math.IsNaN(f) {
			results[i] = nil
		}
	}
	result.value = results

	if result.cacheOpaque == 0 && result.cacheKey != "" {
		// It's OK if we can't store our newly pulled item in
		// the cache, but it's most definitely not OK to stop
		// here because of this.
		select {
		case cacheInputSet <- &result:
		default:
		}
	}
	pi.out <- &result
}

func docProcessor(ch <-chan *processIn) {
	for pi := range ch {
		if time.Now().Before(pi.before) {
			process_docs(pi)
		} else {
			pi.out <- &processOut{"", pi.key, nil, timeoutError, 0}
		}
	}
}

func fetchDocs(dbname string, key int64, infos []*couchstore.DocInfo,
	nextInfo *couchstore.DocInfo, ptrs []string, reds []string,
	filters []string, filtervals []string,
	before time.Time, out chan<- *processOut) {

	i := processIn{"", dbname, key, infos, nextInfo,
		ptrs, reds, before, filters, filtervals, out}

	cacheInput <- &i
}

func runQuery(q *queryIn) {
	if len(q.ptrs) == 0 {
		q.cherr <- fmt.Errorf("At least one pointer is required")
		return
	}

	db, err := dbopen(q.dbname)
	if err != nil {
		log.Printf("Error opening db: %v - %v", q.dbname, err)
		q.cherr <- err
		return
	}
	defer db.Close()

	chunk := int64(time.Duration(q.group) * time.Millisecond)

	infos := []*couchstore.DocInfo{}
	g := int64(0)
	nextg := ""

	err = db.Walk(q.from, func(d *couchstore.Couchstore,
		di *couchstore.DocInfo) error {
		kstr := di.ID()
		var err error
		if q.to != "" && kstr >= q.to {
			err = couchstore.StopIteration
		}

		atomic.AddInt32(&q.totalKeys, 1)

		if kstr >= nextg {
			if len(infos) > 0 {
				atomic.AddInt32(&q.started, 1)
				fetchDocs(q.dbname, g, infos, di,
					q.ptrs, q.reds, q.filters, q.filtervals,
					q.before, q.out)

				infos = make([]*couchstore.DocInfo, 0, len(infos))
			}

			k := parseKey(kstr)
			g = (k / chunk) * chunk
			nextgi := g + chunk
			nextgt := time.Unix(nextgi/1e9, nextgi%1e9).UTC()
			nextg = nextgt.Format(time.RFC3339Nano)
		}
		infos = append(infos, di)

		return err
	})

	if err == nil && len(infos) > 0 {
		atomic.AddInt32(&q.started, 1)
		fetchDocs(q.dbname, g, infos, nil,
			q.ptrs, q.reds, q.filters, q.filtervals,
			q.before, q.out)
	}

	q.cherr <- err
}

func queryExecutor(ch <-chan *queryIn) {
	for q := range ch {
		if time.Now().Before(q.before) {
			runQuery(q)
		} else {
			log.Printf("Timed out query that's %v late",
				time.Since(q.before))
			q.cherr <- timeoutError
		}
	}
}

func executeQuery(dbname, from, to string, group int,
	ptrs []string, reds []string, filters []string, filtervals []string) *queryIn {

	now := time.Now()

	rv := &queryIn{
		dbname:     dbname,
		from:       from,
		to:         to,
		group:      group,
		ptrs:       ptrs,
		reds:       reds,
		start:      now,
		before:     now.Add(*queryTimeout),
		filters:    filters,
		filtervals: filtervals,
		out:        make(chan *processOut),
		cherr:      make(chan error),
	}
	queryInput <- rv
	return rv
}

var processorInput chan *processIn
var queryInput chan *queryIn

func convertTofloat64(in chan ptrval) chan float64 {
	ch := make(chan float64)
	go func() {
		defer close(ch)
		for v := range in {
			if v.included && v.val != nil {
				x, err := strconv.ParseFloat(*v.val, 64)
				if err == nil {
					ch <- x
				}
			}
		}
	}()

	return ch
}

func convertTofloat64Rate(in chan ptrval) chan float64 {
	ch := make(chan float64)
	go func() {
		defer close(ch)
		var prevts int64
		var preval float64

		// First, find a part of the stream that has usable data.
		for v := range in {
			if v.di != nil && v.val != nil {
				x, err := strconv.ParseFloat(*v.val, 64)
				if err == nil {
					prevts = parseKey(v.di.ID())
					preval = x
					break
				}
			}
		}
		// Then emit floats based on deltas from previous values.
		for v := range in {
			if v.di != nil && v.val != nil {
				x, err := strconv.ParseFloat(*v.val, 64)
				if err == nil {
					thists := parseKey(v.di.ID())

					val := ((x - preval) /
						(float64(thists-prevts) / 1e9))

					if !math.IsNaN(val) {
						ch <- val
					}

					prevts = thists
					preval = x
				}
			}
		}
	}()

	return ch
}

var reducers = map[string]Reducer{
	"identity": func(input chan ptrval) interface{} {
		rv := []*string{}
		for s := range input {
			if s.included {
				rv = append(rv, s.val)
			}
		}
		return rv
	},
	"any": func(input chan ptrval) interface{} {
		for v := range input {
			if v.included && v.val != nil {
				return *v.val
			}
		}
		return nil
	},
	"count": func(input chan ptrval) interface{} {
		rv := 0
		for v := range input {
			if v.included && v.val != nil {
				rv++
			}
		}
		return rv
	},
	"sum": func(input chan ptrval) interface{} {
		rv := float64(0)
		for v := range convertTofloat64(input) {
			rv += v
		}
		return rv
	},
	"sumsq": func(input chan ptrval) interface{} {
		rv := float64(0)
		for v := range convertTofloat64(input) {
			rv += (v * v)
		}
		return rv
	},
	"max": func(input chan ptrval) interface{} {
		rv := math.NaN()
		for v := range convertTofloat64(input) {
			if v > rv || math.IsNaN(rv) {
				rv = v
			}
		}
		return rv
	},
	"min": func(input chan ptrval) interface{} {
		rv := math.NaN()
		for v := range convertTofloat64(input) {
			if v < rv || math.IsNaN(rv) {
				rv = v
			}
		}
		return rv
	},
	"avg": func(input chan ptrval) interface{} {
		nums := float64(0)
		sum := float64(0)
		for v := range convertTofloat64(input) {
			nums++
			sum += v
		}
		if nums > 0 {
			return sum / nums
		}
		return math.NaN()
	},
	"c_min": func(input chan ptrval) interface{} {
		rv := math.NaN()
		for v := range convertTofloat64Rate(input) {
			if v < rv || math.IsNaN(rv) {
				rv = v
			}
		}
		return rv
	},
	"c_avg": func(input chan ptrval) interface{} {
		nums := float64(0)
		sum := float64(0)
		for v := range convertTofloat64Rate(input) {
			nums++
			sum += v
		}
		if nums > 0 {
			return sum / nums
		}
		return math.NaN()
	},
	"c_max": func(input chan ptrval) interface{} {
		rv := math.NaN()
		for v := range convertTofloat64Rate(input) {
			if v > rv || math.IsNaN(rv) {
				rv = v
			}
		}
		return rv
	},
}
<|endoftext|>


<commit_before>package tdb

import (
	"bytes"
	"fmt"
	"sort"
	"time"

	"github.com/Knetic/govaluate"
	"github.com/oxtoacart/bytemap"
	"github.com/tecbot/gorocksdb"
)

type query struct {
	table      string
	fields     []string
	filter     *govaluate.EvaluableExpression
	to         time.Time
	toOffset   time.Duration
	from       time.Time
	fromOffset time.Duration
	onValues   func(key bytemap.ByteMap, field string, vals []float64)
}

type QueryStats struct {
	Scanned      int64
	FilterPass   int64
	FilterReject int64
	ReadValue    int64
	DataValid    int64
	InTimeRange  int64
	Runtime      time.Duration
}

func (db *DB) runQuery(q *query) (*QueryStats, error) {
	start := time.Now()
	stats := &QueryStats{}

	if q.from.IsZero() && q.fromOffset >= 0 {
		return stats, fmt.Errorf("Please specify a from or a negative fromOffset")
	}
	if len(q.fields) == 0 {
		return stats, fmt.Errorf("Please specify at least one field")
	}
	t := db.getTable(q.table)
	if t == nil {
		return stats, fmt.Errorf("Unknown table %v", q.table)
	}
	fields := make([][]byte, 0, len(q.fields))
	for _, field := range q.fields {
		fields = append(fields, encodeField(field))
	}
	sort.Sort(lexicographical(fields))
	now := t.clock.Now()
	if q.to.IsZero() {
		q.to = now
		if q.toOffset != 0 {
			q.to = q.to.Add(q.toOffset)
		}
	}
	if q.from.IsZero() {
		q.from = now.Add(q.fromOffset)
	}
	q.to = roundTime(q.to, t.resolution)
	q.from = roundTime(q.from, t.resolution)
	numPeriods := int(q.to.Sub(q.from) / t.resolution)
	log.Tracef("Query will return %d periods for range %v to %v", numPeriods, q.from, q.to)

	ro := gorocksdb.NewDefaultReadOptions()
	// Go ahead and fill the cache
	ro.SetFillCache(true)
	it := t.archiveByKey.NewIterator(ro)
	defer it.Close()

	for _, fieldBytes := range fields {
		for it.Seek(fieldBytes); it.ValidForPrefix(fieldBytes); it.Next() {
			stats.Scanned++
			k := it.Key()
			storedField, key := fieldAndKey(k.Data())

			if q.filter != nil {
				include, err := q.filter.Eval(byteMapParams(key))
				if err != nil {
					k.Free()
					return stats, fmt.Errorf("Unable to apply filter: %v", err)
				}
				inc, ok := include.(bool)
				if !ok {
					k.Free()
					return stats, fmt.Errorf("Filter expression returned something other than a boolean: %v", include)
				}
				if !inc {
					stats.FilterReject++
					continue
				}
				stats.FilterPass++
			}

			v := it.Value()
			stats.ReadValue++
			seq := sequence(v.Data())
			vals := make([]float64, numPeriods)
			if seq.isValid() {
				stats.DataValid++
				seqStart := seq.start()
				if log.IsTraceEnabled() {
					log.Tracef("Sequence starts at %v and has %d periods", seqStart.In(time.UTC), seq.numPeriods())
				}
				includeKey := false
				if !seqStart.Before(q.from) {
					to := q.to
					if to.After(seqStart) {
						to = seqStart
					}
					startOffset := int(seqStart.Sub(to) / t.resolution)
					log.Tracef("Start offset %d", startOffset)
					copyPeriods := seq.numPeriods()
					for i := 0; i+startOffset < copyPeriods && i < numPeriods; i++ {
						includeKey = true
						val := seq.valueAt(i + startOffset)
						log.Tracef("Grabbing value %f", val)
						vals[i] = val
					}
				}
				if includeKey {
					stats.InTimeRange++
					q.onValues(key, storedField, vals)
				}
			}
			k.Free()
			v.Free()
		}
	}

	stats.Runtime = time.Now().Sub(start)
	return stats, nil
}

type byteMapParams bytemap.ByteMap

func (bmp byteMapParams) Get(field string) (interface{}, error) {
	return bytemap.ByteMap(bmp).Get(field), nil
}

type lexicographical [][]byte

func (a lexicographical) Len() int           { return len(a) }
func (a lexicographical) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a lexicographical) Less(i, j int) bool { return bytes.Compare(a[i], a[j]) < 0 }

func keysEqual(a map[string]interface{}, b map[string]interface{}) bool {
	if len(a) != len(b) {
		return false
	}
	for k, v := range a {
		if b[k] != v {
			return false
		}
	}
	return true
}
<commit_msg>Allowing filtering on values that can be nil<commit_after>package tdb

import (
	"bytes"
	"fmt"
	"sort"
	"time"

	"github.com/Knetic/govaluate"
	"github.com/oxtoacart/bytemap"
	"github.com/tecbot/gorocksdb"
)

type query struct {
	table      string
	fields     []string
	filter     *govaluate.EvaluableExpression
	to         time.Time
	toOffset   time.Duration
	from       time.Time
	fromOffset time.Duration
	onValues   func(key bytemap.ByteMap, field string, vals []float64)
}

type QueryStats struct {
	Scanned      int64
	FilterPass   int64
	FilterReject int64
	ReadValue    int64
	DataValid    int64
	InTimeRange  int64
	Runtime      time.Duration
}

func (db *DB) runQuery(q *query) (*QueryStats, error) {
	start := time.Now()
	stats := &QueryStats{}

	if q.from.IsZero() && q.fromOffset >= 0 {
		return stats, fmt.Errorf("Please specify a from or a negative fromOffset")
	}
	if len(q.fields) == 0 {
		return stats, fmt.Errorf("Please specify at least one field")
	}
	t := db.getTable(q.table)
	if t == nil {
		return stats, fmt.Errorf("Unknown table %v", q.table)
	}
	fields := make([][]byte, 0, len(q.fields))
	for _, field := range q.fields {
		fields = append(fields, encodeField(field))
	}
	sort.Sort(lexicographical(fields))
	now := t.clock.Now()
	if q.to.IsZero() {
		q.to = now
		if q.toOffset != 0 {
			q.to = q.to.Add(q.toOffset)
		}
	}
	if q.from.IsZero() {
		q.from = now.Add(q.fromOffset)
	}
	q.to = roundTime(q.to, t.resolution)
	q.from = roundTime(q.from, t.resolution)
	numPeriods := int(q.to.Sub(q.from) / t.resolution)
	log.Tracef("Query will return %d periods for range %v to %v", numPeriods, q.from, q.to)

	ro := gorocksdb.NewDefaultReadOptions()
	// Go ahead and fill the cache
	ro.SetFillCache(true)
	it := t.archiveByKey.NewIterator(ro)
	defer it.Close()

	for _, fieldBytes := range fields {
		for it.Seek(fieldBytes); it.ValidForPrefix(fieldBytes); it.Next() {
			stats.Scanned++
			k := it.Key()
			storedField, key := fieldAndKey(k.Data())

			if q.filter != nil {
				include, err := q.filter.Eval(byteMapParams(key))
				if err != nil {
					k.Free()
					return stats, fmt.Errorf("Unable to apply filter: %v", err)
				}
				inc, ok := include.(bool)
				if !ok {
					k.Free()
					return stats, fmt.Errorf("Filter expression returned something other than a boolean: %v", include)
				}
				if !inc {
					stats.FilterReject++
					continue
				}
				stats.FilterPass++
			}

			v := it.Value()
			stats.ReadValue++
			seq := sequence(v.Data())
			vals := make([]float64, numPeriods)
			if seq.isValid() {
				stats.DataValid++
				seqStart := seq.start()
				if log.IsTraceEnabled() {
					log.Tracef("Sequence starts at %v and has %d periods", seqStart.In(time.UTC), seq.numPeriods())
				}
				includeKey := false
				if !seqStart.Before(q.from) {
					to := q.to
					if to.After(seqStart) {
						to = seqStart
					}
					startOffset := int(seqStart.Sub(to) / t.resolution)
					log.Tracef("Start offset %d", startOffset)
					copyPeriods := seq.numPeriods()
					for i := 0; i+startOffset < copyPeriods && i < numPeriods; i++ {
						includeKey = true
						val := seq.valueAt(i + startOffset)
						log.Tracef("Grabbing value %f", val)
						vals[i] = val
					}
				}
				if includeKey {
					stats.InTimeRange++
					q.onValues(key, storedField, vals)
				}
			}
			k.Free()
			v.Free()
		}
	}

	stats.Runtime = time.Now().Sub(start)
	return stats, nil
}

type byteMapParams bytemap.ByteMap

func (bmp byteMapParams) Get(field string) (interface{}, error) {
	result := bytemap.ByteMap(bmp).Get(field)
	if result == nil {
		return "", nil
	}
	return result, nil
}

type lexicographical [][]byte

func (a lexicographical) Len() int           { return len(a) }
func (a lexicographical) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a lexicographical) Less(i, j int) bool { return bytes.Compare(a[i], a[j]) < 0 }

func keysEqual(a map[string]interface{}, b map[string]interface{}) bool {
	if len(a) != len(b) {
		return false
	}
	for k, v := range a {
		if b[k] != v {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"time"
)

var quests List

func randQuest() string {
	num := rand.Int31n(int32(len(quests.Quests)))
	return fmt.Sprintf("%s\t%s", quests.Quests[num].Rank, quests.Quests[num].Name)
}

func readQuests(name string) error {
	file, err := os.Open(name)
	if err != nil {
		return err
	}
	content, err := ioutil.ReadAll(file)
	json.Unmarshal(content, &quests)
	return nil
}

func root(rw http.ResponseWriter, req *http.Request) {
	output := randQuest()
	io.WriteString(rw, output+"\n")
	fmt.Println(output)
}

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

func main() {
	input := flag.String("file", "quests.json", "Quests json location")
	httpPtr := flag.Bool("http", false, "Open http server")
	ip := flag.String("ip", "0.0.0.0", "http server ip")
	port := flag.Int("port", 8080, "http server port")
	flag.Parse()
	err := readQuests(*input)
	if err != nil {
		fmt.Println("Read Quest Error!!")
		return
	}
	if *httpPtr == false {
		fmt.Println(randQuest())
		return
	}
	http.Handle("/root/", http.HandlerFunc(root))
	err = http.ListenAndServe(fmt.Sprintf("%s:%d", *ip, *port), nil)
	if err != nil {
		log.Fatal("ListenAndServe", err)
	}
}
<commit_msg>Alter: root handle uri modify<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"math/rand"
	"net/http"
	"os"
	"time"
)

var quests List

func randQuest() string {
	num := rand.Int31n(int32(len(quests.Quests)))
	return fmt.Sprintf("%s\t%s", quests.Quests[num].Rank, quests.Quests[num].Name)
}

func readQuests(name string) error {
	file, err := os.Open(name)
	if err != nil {
		return err
	}
	content, err := ioutil.ReadAll(file)
	json.Unmarshal(content, &quests)
	return nil
}

func root(rw http.ResponseWriter, req *http.Request) {
	output := randQuest()
	io.WriteString(rw, output+"\n")
	fmt.Println(output)
}

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

func main() {
	input := flag.String("file", "quests.json", "Quests json location")
	httpPtr := flag.Bool("http", false, "Open http server")
	ip := flag.String("ip", "0.0.0.0", "http server ip")
	port := flag.Int("port", 8080, "http server port")
	flag.Parse()
	err := readQuests(*input)
	if err != nil {
		fmt.Println("Read Quest Error!!")
		return
	}
	if *httpPtr == false {
		fmt.Println(randQuest())
		return
	}
	http.Handle("/", http.HandlerFunc(root))
	fmt.Printf("Serve MH4G Quest Random Generator at %s:%d\n", *ip, *port)
	err = http.ListenAndServe(fmt.Sprintf("%s:%d", *ip, *port), nil)
	if err != nil {
		log.Fatal("ListenAndServe", err)
	}
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package tests_test

import (
	"context"
	"fmt"
	"strings"
	"time"

	expect "github.com/google/goexpect"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"

	"kubevirt.io/kubevirt/tests/util"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/kubevirt/tests"
	"kubevirt.io/kubevirt/tests/console"
	cd "kubevirt.io/kubevirt/tests/containerdisk"
	"kubevirt.io/kubevirt/tests/libvmi"
)

var _ = Describe("[rfe_id:588][crit:medium][vendor:cnv-qe@redhat.com][level:component][sig-compute]ContainerDisk", func() {

	var virtClient kubecli.KubevirtClient
	var err error

	BeforeEach(func() {
		virtClient, err = kubecli.GetKubevirtClient()
		util.PanicOnError(err)
	})

	AfterEach(func() {
		tests.BeforeTestCleanup()
	})

	verifyContainerDiskVMI := func(vmi *v1.VirtualMachineInstance, obj runtime.Object) {
		_, ok := obj.(*v1.VirtualMachineInstance)
		Expect(ok).To(BeTrue(), "Object is not of type *v1.VirtualMachineInstance")
		tests.WaitForSuccessfulVMIStart(obj)

		// Verify Registry Disks are Online
		pods, err := virtClient.CoreV1().Pods(util.NamespaceTestDefault).List(context.Background(), tests.UnfinishedVMIPodSelector(vmi))
		Expect(err).To(BeNil())

		By("Checking the number of VirtualMachineInstance disks")
		disksFound := 0
		for _, pod := range pods.Items {
			if pod.ObjectMeta.DeletionTimestamp != nil {
				continue
			}
			for _, containerStatus := range pod.Status.ContainerStatuses {
				if strings.HasPrefix(containerStatus.Name, "volume") == false {
					// only check readiness of disk containers
					continue
				}
				disksFound++
			}
			break
		}
		Expect(disksFound).To(Equal(1))
	}

	table.DescribeTable("should", func(image string, policy k8sv1.PullPolicy, expectedPolicy k8sv1.PullPolicy) {
		vmi := tests.NewRandomVMIWithEphemeralDisk(image)
		vmi.Spec.Volumes[0].ContainerDisk.ImagePullPolicy = policy
		vmi = tests.RunVMIAndExpectScheduling(vmi, 60)
		Expect(vmi.Spec.Volumes[0].ContainerDisk.ImagePullPolicy).To(Equal(expectedPolicy))
		pod := libvmi.GetPodByVirtualMachineInstance(vmi, vmi.Namespace)
		container := tests.GetContainerDiskContainerOfPod(pod, vmi.Spec.Volumes[0].Name)
		Expect(container.ImagePullPolicy).To(Equal(expectedPolicy))
	},
		table.Entry("[test_id:3246]generate and set Always pull policy", "test", k8sv1.PullPolicy(""), k8sv1.PullAlways),
		table.Entry("[test_id:3247]generate and set Always pull policy", "test:latest", k8sv1.PullPolicy(""), k8sv1.PullAlways),
		table.Entry("[test_id:3248]generate and set IfNotPresent pull policy", "test@sha256:9c2b78e11c25b3fd0b24b0ed684a112052dff03eee4ca4bdcc4f3168f9a14396", k8sv1.PullPolicy(""), k8sv1.PullIfNotPresent),
		table.Entry("[test_id:3249]pass through Never pull policy to the pod", "test@sha256:9c2b78e11c25b3fd0b24b0ed684a112052dff03eee4ca4bdcc4f3168f9a14396", k8sv1.PullNever, k8sv1.PullNever),
		table.Entry("[test_id:3250]pass through IfNotPresent pull policy to the pod", "test:latest", k8sv1.PullIfNotPresent, k8sv1.PullIfNotPresent),
	)

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting and stopping the same VirtualMachineInstance", func() {
		Context("with ephemeral registry disk", func() {
			It("[test_id:1463][Conformance] should success multiple times", func() {
				vmi := tests.NewRandomVMIWithEphemeralDiskAndUserdata(cd.ContainerDiskFor(cd.ContainerDiskCirros), "#!/bin/bash\necho 'hello'\n")
				num := 2
				for i := 0; i < num; i++ {
					By("Starting the VirtualMachineInstance")
					obj, err := virtClient.RestClient().Post().Resource("virtualmachineinstances").Namespace(util.NamespaceTestDefault).Body(vmi).Do(context.Background()).Get()
					Expect(err).To(BeNil())
					tests.WaitForSuccessfulVMIStart(obj)

					By("Stopping the VirtualMachineInstance")
					_, err = virtClient.RestClient().Delete().Resource("virtualmachineinstances").Namespace(vmi.GetObjectMeta().GetNamespace()).Name(vmi.GetObjectMeta().GetName()).Do(context.Background()).Get()
					Expect(err).To(BeNil())
					By("Waiting until the VirtualMachineInstance is gone")
					tests.WaitForVirtualMachineToDisappearWithTimeout(vmi, 120)
				}
			})
		})
	})

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting a VirtualMachineInstance", func() {
		Context("with ephemeral registry disk", func() {
			It("[test_id:1464]should not modify the spec on status update", func() {
				vmi := tests.NewRandomVMIWithEphemeralDiskAndUserdata(cd.ContainerDiskFor(cd.ContainerDiskCirros), "#!/bin/bash\necho 'hello'\n")
				v1.SetObjectDefaults_VirtualMachineInstance(vmi)

				By("Starting the VirtualMachineInstance")
				vmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
				Expect(err).To(BeNil())
				tests.WaitForSuccessfulVMIStart(vmi)
				startedVMI, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Get(vmi.ObjectMeta.Name, &metav1.GetOptions{})
				Expect(err).To(BeNil())
				By("Checking that the VirtualMachineInstance spec did not change")
				Expect(startedVMI.Spec).To(Equal(vmi.Spec))
			})
		})
		Context("[Serial]should obey the disk verification limits in the KubeVirt CR", func() {
			It("[test_id:7182]disk verification should fail when the memory limit is too low", func() {
				By("Reducing the diskVerificaton memory usage limit")
				kv := util.GetCurrentKv(virtClient)
				kv.Spec.Configuration.DeveloperConfiguration = &v1.DeveloperConfiguration{
					DiskVerification: &v1.DiskVerification{
						MemoryLimit: resource.NewScaledQuantity(42, resource.Kilo),
					},
				}
				tests.UpdateKubeVirtConfigValueAndWait(kv.Spec.Configuration)
				Expect(err).To(BeNil())

				By("Starting the VirtualMachineInstance")
				vmi := tests.NewRandomVMIWithEphemeralDisk(cd.ContainerDiskFor(cd.ContainerDiskCirros))
				_, err = virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
				Expect(err).To(BeNil())
				By("Checking that the VMI failed")
				Eventually(func() bool {
					vmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Get(vmi.Name, &metav1.GetOptions{})
					Expect(err).ToNot(HaveOccurred())
					for _, condition := range vmi.Status.Conditions {
						if condition.Type == v1.VirtualMachineInstanceSynchronized && condition.Status == k8sv1.ConditionFalse {
							return strings.Contains(condition.Message, "failed to invoke qemu-img")
						}
					}
					return false
				}, 3*time.Minute, 2*time.Second).Should(BeTrue())
			})
		})
	})

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting multiple VMIs", func() {
		Context("with ephemeral registry disk", func() {
			It("[test_id:1465]should success", func() {
				num := 5
				vmis := make([]*v1.VirtualMachineInstance, 0, num)
				objs := make([]runtime.Object, 0, num)
				for i := 0; i < num; i++ {
					vmi := tests.NewRandomVMIWithEphemeralDiskAndUserdata(cd.ContainerDiskFor(cd.ContainerDiskCirros), "#!/bin/bash\necho 'hello'\n")
					// FIXME if we give too much ram, the vmis really boot and eat all our memory (cache?)
					vmi.Spec.Domain.Resources.Requests[k8sv1.ResourceMemory] = resource.MustParse("1M")
					obj := tests.RunVMI(vmi, 10)
					vmis = append(vmis, vmi)
					objs = append(objs, obj)
				}

				for idx, vmi := range vmis {
					verifyContainerDiskVMI(vmi, objs[idx])
				}
			}) // Timeout is long because this test involves multiple parallel VirtualMachineInstance launches.
		})
	})

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting from custom image location", func() {
		Context("with disk at /custom-disk/downloaded", func() {
			It("[test_id:1466]should boot normally", func() {
				vmi := tests.NewRandomVMIWithEphemeralDiskAndUserdata(cd.ContainerDiskFor(cd.ContainerDiskCirrosCustomLocation), "#!/bin/bash\necho 'hello'\n")
				for ind, volume := range vmi.Spec.Volumes {
					if volume.ContainerDisk != nil {
						vmi.Spec.Volumes[ind].ContainerDisk.Path = "/custom-disk/downloaded"
					}
				}
				By("Starting the VirtualMachineInstance")
				obj, err := virtClient.RestClient().Post().Resource("virtualmachineinstances").Namespace(util.NamespaceTestDefault).Body(vmi).Do(context.Background()).Get()
				Expect(err).To(BeNil())
				tests.WaitForSuccessfulVMIStart(obj)
			})
		})

	})

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting with virtio-win", func() {
		Context("with virtio-win as secondary disk", func() {
			It("[test_id:1467]should boot and have the virtio as sata CDROM", func() {
				vmi := tests.NewRandomVMIWithEphemeralDisk(cd.ContainerDiskFor(cd.ContainerDiskAlpine))
				tests.AddEphemeralCdrom(vmi, "disk4", "sata", cd.ContainerDiskFor(cd.ContainerDiskVirtio))

				By("Starting the VirtualMachineInstance")
				obj, err := virtClient.RestClient().Post().Resource("virtualmachineinstances").Namespace(util.NamespaceTestDefault).Body(vmi).Do(context.Background()).Get()
				Expect(err).To(BeNil(), "expected vmi to start with no problem")
				tests.WaitForSuccessfulVMIStart(obj)

				By("Checking whether the second disk really contains virtio drivers")
				Expect(console.LoginToAlpine(vmi)).To(Succeed(), "expected alpine to login properly")

				Expect(console.SafeExpectBatch(vmi, []expect.Batcher{
					// mount virtio cdrom and check files are there
					&expect.BSnd{S: "mount -t iso9600 /dev/cdrom\n"},
					&expect.BExp{R: console.PromptExpression},
					&expect.BSnd{S: "echo $?\n"},
					&expect.BExp{R: console.RetValue("0")},
					&expect.BSnd{S: "cd /media/cdrom\n"},
					&expect.BExp{R: console.PromptExpression},
					&expect.BSnd{S: "ls virtio-win_license.txt guest-agent\n"},
					&expect.BExp{R: console.PromptExpression},
					&expect.BSnd{S: "echo $?\n"},
					&expect.BExp{R: console.RetValue("0")},
				}, 200)).To(Succeed(), "expected virtio files to be mounted properly")
			})
		})
	})

	Describe("[rfe_id:4052][crit:high][arm64][vendor:cnv-qe@redhat.com][level:component]VMI disk permissions", func() {
		Context("with ephemeral registry disk", func() {
			It("[test_id:4299]should not have world write permissions", func() {
				vmi := tests.NewRandomVMIWithEphemeralDisk(cd.ContainerDiskFor(cd.ContainerDiskAlpine))

				By("Starting a New VMI")
				vmi, err = virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
				Expect(err).ToNot(HaveOccurred())
				tests.WaitForSuccessfulVMIStart(vmi)

				By("Ensuring VMI is running by logging in")
				tests.WaitUntilVMIReady(vmi, console.LoginToAlpine)

				By("Fetching virt-launcher Pod")
				pod := libvmi.GetPodByVirtualMachineInstance(vmi, util.NamespaceTestDefault)

				writableImagePath := fmt.Sprintf("/var/run/kubevirt-ephemeral-disks/disk-data/%v/disk.qcow2", vmi.Spec.Domain.Devices.Disks[0].Name)

				writableImageOctalMode, err := tests.ExecuteCommandOnPod(
					virtClient,
					pod,
					"compute",
					[]string{"/usr/bin/bash", "-c", fmt.Sprintf("stat -c %%a %s", writableImagePath)},
				)
				Expect(err).ToNot(HaveOccurred())

				By("Checking the writable Image Octal mode")
				Expect(strings.Trim(writableImageOctalMode, "\n")).To(Equal("640"), "Octal Mode of writable Image should be 640")

				readonlyImageOctalMode, err := tests.ExecuteCommandOnPod(
					virtClient,
					pod,
					"compute",
					[]string{"/usr/bin/bash", "-c", "stat -c %a /var/run/kubevirt/container-disks/disk_0.img"},
				)
				Expect(err).ToNot(HaveOccurred())

				By("Checking the read-only Image Octal mode")
				Expect(strings.Trim(readonlyImageOctalMode, "\n")).To(Equal("440"), "Octal Mode of read-only Image should be 440")

			})
		})
	})
})
<commit_msg>container_disk_test: reuse libvmi.New* and RunVMIAndExpectLaunch<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package tests_test

import (
	"context"
	"fmt"
	"strings"
	"time"

	expect "github.com/google/goexpect"
	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	k8sv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"

	"kubevirt.io/kubevirt/tests/util"

	v1 "kubevirt.io/api/core/v1"
	"kubevirt.io/client-go/kubecli"
	"kubevirt.io/kubevirt/tests"
	"kubevirt.io/kubevirt/tests/console"
	cd "kubevirt.io/kubevirt/tests/containerdisk"
	"kubevirt.io/kubevirt/tests/libvmi"
)

var _ = Describe("[rfe_id:588][crit:medium][vendor:cnv-qe@redhat.com][level:component][sig-compute]ContainerDisk", func() {

	var virtClient kubecli.KubevirtClient
	var err error

	BeforeEach(func() {
		virtClient, err = kubecli.GetKubevirtClient()
		util.PanicOnError(err)
	})

	AfterEach(func() {
		tests.BeforeTestCleanup()
	})

	verifyContainerDiskVMI := func(vmi *v1.VirtualMachineInstance, obj runtime.Object) {
		_, ok := obj.(*v1.VirtualMachineInstance)
		Expect(ok).To(BeTrue(), "Object is not of type *v1.VirtualMachineInstance")
		tests.WaitForSuccessfulVMIStart(obj)

		// Verify Registry Disks are Online
		pods, err := virtClient.CoreV1().Pods(util.NamespaceTestDefault).List(context.Background(), tests.UnfinishedVMIPodSelector(vmi))
		Expect(err).To(BeNil())

		By("Checking the number of VirtualMachineInstance disks")
		disksFound := 0
		for _, pod := range pods.Items {
			if pod.ObjectMeta.DeletionTimestamp != nil {
				continue
			}
			for _, containerStatus := range pod.Status.ContainerStatuses {
				if strings.HasPrefix(containerStatus.Name, "volume") == false {
					// only check readiness of disk containers
					continue
				}
				disksFound++
			}
			break
		}
		Expect(disksFound).To(Equal(1))
	}

	table.DescribeTable("should", func(image string, policy k8sv1.PullPolicy, expectedPolicy k8sv1.PullPolicy) {
		vmi := tests.NewRandomVMIWithEphemeralDisk(image)
		vmi.Spec.Volumes[0].ContainerDisk.ImagePullPolicy = policy
		vmi = tests.RunVMIAndExpectScheduling(vmi, 60)
		Expect(vmi.Spec.Volumes[0].ContainerDisk.ImagePullPolicy).To(Equal(expectedPolicy))
		pod := libvmi.GetPodByVirtualMachineInstance(vmi, vmi.Namespace)
		container := tests.GetContainerDiskContainerOfPod(pod, vmi.Spec.Volumes[0].Name)
		Expect(container.ImagePullPolicy).To(Equal(expectedPolicy))
	},
		table.Entry("[test_id:3246]generate and set Always pull policy", "test", k8sv1.PullPolicy(""), k8sv1.PullAlways),
		table.Entry("[test_id:3247]generate and set Always pull policy", "test:latest", k8sv1.PullPolicy(""), k8sv1.PullAlways),
		table.Entry("[test_id:3248]generate and set IfNotPresent pull policy", "test@sha256:9c2b78e11c25b3fd0b24b0ed684a112052dff03eee4ca4bdcc4f3168f9a14396", k8sv1.PullPolicy(""), k8sv1.PullIfNotPresent),
		table.Entry("[test_id:3249]pass through Never pull policy to the pod", "test@sha256:9c2b78e11c25b3fd0b24b0ed684a112052dff03eee4ca4bdcc4f3168f9a14396", k8sv1.PullNever, k8sv1.PullNever),
		table.Entry("[test_id:3250]pass through IfNotPresent pull policy to the pod", "test:latest", k8sv1.PullIfNotPresent, k8sv1.PullIfNotPresent),
	)

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting and stopping the same VirtualMachineInstance", func() {
		Context("with ephemeral registry disk", func() {
			It("[test_id:1463][Conformance] should success multiple times", func() {
				vmi := tests.NewRandomVMIWithEphemeralDiskAndUserdata(cd.ContainerDiskFor(cd.ContainerDiskCirros), "#!/bin/bash\necho 'hello'\n")
				num := 2
				for i := 0; i < num; i++ {
					By("Starting the VirtualMachineInstance")
					obj, err := virtClient.RestClient().Post().Resource("virtualmachineinstances").Namespace(util.NamespaceTestDefault).Body(vmi).Do(context.Background()).Get()
					Expect(err).To(BeNil())
					tests.WaitForSuccessfulVMIStart(obj)

					By("Stopping the VirtualMachineInstance")
					_, err = virtClient.RestClient().Delete().Resource("virtualmachineinstances").Namespace(vmi.GetObjectMeta().GetNamespace()).Name(vmi.GetObjectMeta().GetName()).Do(context.Background()).Get()
					Expect(err).To(BeNil())
					By("Waiting until the VirtualMachineInstance is gone")
					tests.WaitForVirtualMachineToDisappearWithTimeout(vmi, 120)
				}
			})
		})
	})

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting a VirtualMachineInstance", func() {
		Context("with ephemeral registry disk", func() {
			It("[test_id:1464]should not modify the spec on status update", func() {
				vmi := libvmi.NewCirros()
				v1.SetObjectDefaults_VirtualMachineInstance(vmi)

				By("Starting the VirtualMachineInstance")
				vmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
				Expect(err).To(BeNil())
				tests.WaitForSuccessfulVMIStart(vmi)
				startedVMI, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Get(vmi.ObjectMeta.Name, &metav1.GetOptions{})
				Expect(err).To(BeNil())
				By("Checking that the VirtualMachineInstance spec did not change")
				Expect(startedVMI.Spec).To(Equal(vmi.Spec))
			})
		})
		Context("[Serial]should obey the disk verification limits in the KubeVirt CR", func() {
			It("[test_id:7182]disk verification should fail when the memory limit is too low", func() {
				By("Reducing the diskVerificaton memory usage limit")
				kv := util.GetCurrentKv(virtClient)
				kv.Spec.Configuration.DeveloperConfiguration = &v1.DeveloperConfiguration{
					DiskVerification: &v1.DiskVerification{
						MemoryLimit: resource.NewScaledQuantity(42, resource.Kilo),
					},
				}
				tests.UpdateKubeVirtConfigValueAndWait(kv.Spec.Configuration)
				Expect(err).To(BeNil())

				By("Starting the VirtualMachineInstance")
				vmi := libvmi.NewCirros()
				_, err = virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Create(vmi)
				Expect(err).To(BeNil())
				By("Checking that the VMI failed")
				Eventually(func() bool {
					vmi, err := virtClient.VirtualMachineInstance(util.NamespaceTestDefault).Get(vmi.Name, &metav1.GetOptions{})
					Expect(err).ToNot(HaveOccurred())
					for _, condition := range vmi.Status.Conditions {
						if condition.Type == v1.VirtualMachineInstanceSynchronized && condition.Status == k8sv1.ConditionFalse {
							return strings.Contains(condition.Message, "failed to invoke qemu-img")
						}
					}
					return false
				}, 3*time.Minute, 2*time.Second).Should(BeTrue())
			})
		})
	})

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting multiple VMIs", func() {
		Context("with ephemeral registry disk", func() {
			It("[test_id:1465]should success", func() {
				num := 5
				vmis := make([]*v1.VirtualMachineInstance, 0, num)
				objs := make([]runtime.Object, 0, num)
				for i := 0; i < num; i++ {
					vmi := tests.NewRandomVMIWithEphemeralDiskAndUserdata(cd.ContainerDiskFor(cd.ContainerDiskCirros), "#!/bin/bash\necho 'hello'\n")
					// FIXME if we give too much ram, the vmis really boot and eat all our memory (cache?)
					vmi.Spec.Domain.Resources.Requests[k8sv1.ResourceMemory] = resource.MustParse("1M")
					obj := tests.RunVMI(vmi, 10)
					vmis = append(vmis, vmi)
					objs = append(objs, obj)
				}

				for idx, vmi := range vmis {
					verifyContainerDiskVMI(vmi, objs[idx])
				}
			}) // Timeout is long because this test involves multiple parallel VirtualMachineInstance launches.
		})
	})

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting from custom image location", func() {
		Context("with disk at /custom-disk/downloaded", func() {
			It("[test_id:1466]should boot normally", func() {
				vmi := tests.NewRandomVMIWithEphemeralDiskAndUserdata(cd.ContainerDiskFor(cd.ContainerDiskCirrosCustomLocation), "#!/bin/bash\necho 'hello'\n")
				for ind, volume := range vmi.Spec.Volumes {
					if volume.ContainerDisk != nil {
						vmi.Spec.Volumes[ind].ContainerDisk.Path = "/custom-disk/downloaded"
					}
				}
				By("Starting the VirtualMachineInstance")
				vmi = tests.RunVMIAndExpectLaunch(vmi, 60)
			})
		})

	})

	Describe("[rfe_id:273][crit:medium][vendor:cnv-qe@redhat.com][level:component]Starting with virtio-win", func() {
		Context("with virtio-win as secondary disk", func() {
			It("[test_id:1467]should boot and have the virtio as sata CDROM", func() {
				vmi := libvmi.NewAlpine()
				tests.AddEphemeralCdrom(vmi, "disk4", "sata", cd.ContainerDiskFor(cd.ContainerDiskVirtio))
				vmi = tests.RunVMIAndExpectLaunch(vmi, 60)

				By("Checking whether the second disk really contains virtio drivers")
				Expect(console.LoginToAlpine(vmi)).To(Succeed(), "expected alpine to login properly")

				Expect(console.SafeExpectBatch(vmi, []expect.Batcher{
					// mount virtio cdrom and check files are there
					&expect.BSnd{S: "mount -t iso9600 /dev/cdrom\n"},
					&expect.BExp{R: console.PromptExpression},
					&expect.BSnd{S: "echo $?\n"},
					&expect.BExp{R: console.RetValue("0")},
					&expect.BSnd{S: "cd /media/cdrom\n"},
					&expect.BExp{R: console.PromptExpression},
					&expect.BSnd{S: "ls virtio-win_license.txt guest-agent\n"},
					&expect.BExp{R: console.PromptExpression},
					&expect.BSnd{S: "echo $?\n"},
					&expect.BExp{R: console.RetValue("0")},
				}, 200)).To(Succeed(), "expected virtio files to be mounted properly")
			})
		})
	})

	Describe("[rfe_id:4052][crit:high][arm64][vendor:cnv-qe@redhat.com][level:component]VMI disk permissions", func() {
		Context("with ephemeral registry disk", func() {
			It("[test_id:4299]should not have world write permissions", func() {
				vmi := libvmi.NewAlpine()
				vmi = tests.RunVMIAndExpectLaunch(vmi, 60)

				By("Ensuring VMI is running by logging in")
				tests.WaitUntilVMIReady(vmi, console.LoginToAlpine)

				By("Fetching virt-launcher Pod")
				pod := libvmi.GetPodByVirtualMachineInstance(vmi, util.NamespaceTestDefault)

				writableImagePath := fmt.Sprintf("/var/run/kubevirt-ephemeral-disks/disk-data/%v/disk.qcow2", vmi.Spec.Domain.Devices.Disks[0].Name)

				writableImageOctalMode, err := tests.ExecuteCommandOnPod(
					virtClient,
					pod,
					"compute",
					[]string{"/usr/bin/bash", "-c", fmt.Sprintf("stat -c %%a %s", writableImagePath)},
				)
				Expect(err).ToNot(HaveOccurred())

				By("Checking the writable Image Octal mode")
				Expect(strings.Trim(writableImageOctalMode, "\n")).To(Equal("640"), "Octal Mode of writable Image should be 640")

				readonlyImageOctalMode, err := tests.ExecuteCommandOnPod(
					virtClient,
					pod,
					"compute",
					[]string{"/usr/bin/bash", "-c", "stat -c %a /var/run/kubevirt/container-disks/disk_0.img"},
				)
				Expect(err).ToNot(HaveOccurred())

				By("Checking the read-only Image Octal mode")
				Expect(strings.Trim(readonlyImageOctalMode, "\n")).To(Equal("440"), "Octal Mode of read-only Image should be 440")

			})
		})
	})
})
<|endoftext|>


<commit_before>package ratio

import (
	"io"
	"time"
)

type values struct {
	n   int
	err error
}

type operation struct {
	p      []byte
	values chan values
}

type rateLimiter struct {
	stop      chan bool
	worker    chan *operation
	action    func(p []byte) (int, error)
	limit     int
	remaining int
	written   int
	op        *operation
}

func (rl *rateLimiter) reset() {
	rl.remaining = rl.limit
}

func (rl *rateLimiter) run(duration time.Duration) {
	for {
		select {
		case <-rl.stop:
			rl.close()
			return
		case <-time.Tick(duration):
			rl.reset()
			rl.write()
		case op := <-rl.worker:
			rl.op = op
			rl.write()
		}
	}
}

func (rl *rateLimiter) write() {
	if rl.op == nil {
		return
	}
	p := rl.op.p
	if rl.remaining < len(p) {
		p = rl.op.p[:rl.remaining]
	}
	rl.remaining -= len(p)

	n, err := rl.action(p)
	rl.op.p = rl.op.p[n:]
	rl.written += n

	if len(rl.op.p) == 0 || err != nil {
		rl.op.values <- values{rl.written, err}
		rl.op = nil
	}
}

func (rl *rateLimiter) record(p []byte) (int, error) {
	op := &operation{
		p:      p,
		values: make(chan values, 1),
	}
	select {
	case <-rl.stop:
		return 0, io.EOF
	case rl.worker <- op:
		val := <-op.values
		return val.n, val.err
	}
}

func (rl *rateLimiter) close() {
	if rl.op != nil {
		rl.op.values <- values{rl.written, io.EOF}
	}
}

// Returns a rate limited io.Writer
func NewRateWriter(w io.Writer, size int, duration time.Duration) io.WriteCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return w.Write(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

// Returns a rate limited io.Reader
func NewRateReader(r io.Reader, size int, duration time.Duration) io.ReadCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return r.Read(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

func (rl *rateLimiter) Read(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Write(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Close() error {
	select {
	case <-rl.stop:
		return io.EOF
	default:
	}
	close(rl.stop)
	return nil
}
<commit_msg>improve doc<commit_after>// Rate limited io
package ratio

import (
	"io"
	"time"
)

type values struct {
	n   int
	err error
}

type operation struct {
	p      []byte
	values chan values
}

type rateLimiter struct {
	stop      chan bool
	worker    chan *operation
	action    func(p []byte) (int, error)
	limit     int
	remaining int
	written   int
	op        *operation
}

func (rl *rateLimiter) reset() {
	rl.remaining = rl.limit
}

func (rl *rateLimiter) run(duration time.Duration) {
	for {
		select {
		case <-rl.stop:
			rl.close()
			return
		case <-time.Tick(duration):
			rl.reset()
			rl.write()
		case op := <-rl.worker:
			rl.op = op
			rl.write()
		}
	}
}

func (rl *rateLimiter) write() {
	if rl.op == nil {
		return
	}
	p := rl.op.p
	if rl.remaining < len(p) {
		p = rl.op.p[:rl.remaining]
	}
	rl.remaining -= len(p)

	n, err := rl.action(p)
	rl.op.p = rl.op.p[n:]
	rl.written += n

	if len(rl.op.p) == 0 || err != nil {
		rl.op.values <- values{rl.written, err}
		rl.op = nil
	}
}

func (rl *rateLimiter) record(p []byte) (int, error) {
	op := &operation{
		p:      p,
		values: make(chan values, 1),
	}
	select {
	case <-rl.stop:
		return 0, io.EOF
	case rl.worker <- op:
		val := <-op.values
		return val.n, val.err
	}
}

func (rl *rateLimiter) close() {
	if rl.op != nil {
		rl.op.values <- values{rl.written, io.EOF}
	}
}

// Returns a rate limited io.Writer, allowing to write up to size bytes per duration.
func NewRateWriter(w io.Writer, size int, duration time.Duration) io.WriteCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return w.Write(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

// Returns a rate limited io.Reader, allowing to read up to size bytes per duration.
func NewRateReader(r io.Reader, size int, duration time.Duration) io.ReadCloser {
	rl := &rateLimiter{
		action: func(p []byte) (int, error) {
			return r.Read(p)
		},
		limit:     size,
		remaining: size,
		stop:      make(chan bool),
		worker:    make(chan *operation),
	}
	go rl.run(duration)
	return rl
}

func (rl *rateLimiter) Read(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Write(p []byte) (int, error) {
	return rl.record(p)
}

func (rl *rateLimiter) Close() error {
	select {
	case <-rl.stop:
		return io.EOF
	default:
	}
	close(rl.stop)
	return nil
}
<|endoftext|>