<commit_before>package main

import (
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"os"
	"sync/atomic"
	"time"

	"github.com/gorilla/context"

	"github.com/getlantern/measured"

	"github.com/getlantern/http-proxy-extensions/devicefilter"
	"github.com/getlantern/http-proxy-extensions/mimic"
	"github.com/getlantern/http-proxy-extensions/profilter"
	"github.com/getlantern/http-proxy-extensions/tokenfilter"
	"github.com/getlantern/http-proxy/commonfilter"
	"github.com/getlantern/http-proxy/forward"
	"github.com/getlantern/http-proxy/httpconnect"
	"github.com/getlantern/http-proxy/utils"
)

type Server struct {
	firstHandler http.Handler
	httpServer   http.Server
	tls          bool

	listener net.Listener

	maxConns uint64
	numConns uint64

	idleTimeout time.Duration
}

func NewServer(token string, maxConns uint64, idleTimeout time.Duration, enableFilters bool, logLevel utils.LogLevel) *Server {
	stdWriter := io.Writer(os.Stdout)

	if maxConns == 0 {
		maxConns = math.MaxInt64
	}

	// The following middleware architecture can be seen as a chain of
	// filters that is run from last to first.
	// Don't forget to check Oxy and Gorilla's handlers for middleware.

	// Handles Direct Proxying
	forwardHandler, _ := forward.New(
		nil,
		forward.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		forward.IdleTimeoutSetter(idleTimeout),
	)

	// Handles HTTP CONNECT
	connectHandler, _ := httpconnect.New(
		forwardHandler,
		httpconnect.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		httpconnect.IdleTimeoutSetter(idleTimeout),
	)

	// Catches any request before reaching the CONNECT middleware or
	// the forwarder
	commonFilter, _ := commonfilter.New(
		connectHandler,
		commonfilter.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
	)

	var firstHandler http.Handler
	if !enableFilters {
		firstHandler = commonFilter
	} else {
		// Identifies Lantern Pro users (currently NOOP)
		lanternPro, _ := profilter.New(
			commonFilter,
			profilter.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		)
		// Returns a 404 to requests without the proper token.  Removes the
		// header before continuing.
		tokenFilter, _ := tokenfilter.New(
			lanternPro,
			tokenfilter.TokenSetter(token),
			tokenfilter.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		)
		// Extracts the user ID and attaches the matching client to the request
		// context.  Returns a 404 to requests without the UID.  Removes the
		// header before continuing.
		deviceFilter, _ := devicefilter.New(
			tokenFilter,
			devicefilter.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		)
		firstHandler = deviceFilter
	}

	server := &Server{
		firstHandler: firstHandler,
		maxConns:     maxConns,
		idleTimeout:  idleTimeout,
	}
	return server
}

func (s *Server) ServeHTTP(addr string, chListenOn *chan string) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	s.tls = false
	fmt.Printf("Listen http on %s\n", addr)
	return s.doServe(listener, chListenOn)
}

func (s *Server) ServeHTTPS(addr, keyfile, certfile string, chListenOn *chan string) error {
	listener, err := listenTLS(addr, keyfile, certfile)
	if err != nil {
		return err
	}
	s.tls = true
	fmt.Printf("Listen http on %s\n", addr)
	return s.doServe(listener, chListenOn)
}

func (s *Server) doServe(listener net.Listener, chListenOn *chan string) error {
	// A dirty trick to associate a connection with the http.Request it
	// contains. In "net/http/server.go", handler will be called
	// immediately after ConnState changed to StateActive, so it's safe to
	// loop through all elements in a channel to find a match remote addr.
	q := make(chan net.Conn, 10)

	proxy := http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			for c := range q {
				if c.RemoteAddr().String() == req.RemoteAddr {
					context.Set(req, "conn", c)
					break
				} else {
					q <- c
				}
			}
			s.firstHandler.ServeHTTP(w, req)
		})

	limListener := newLimitedListener(listener, &s.numConns, s.idleTimeout)
	if *enableReports {
		mListener := measured.Listener(limListener, 30*time.Second)
		s.listener = mListener
	} else {
		s.listener = limListener
	}

	s.httpServer = http.Server{Handler: proxy,
		ConnState: func(c net.Conn, state http.ConnState) {
			if state == http.StateActive {
				select {
				case q <- c:
				default:
					fmt.Print("Oops! the connection queue is full!\n")
				}
			}

			if atomic.LoadUint64(&s.numConns) >= s.maxConns {
				limListener.Stop()
			} else if limListener.IsStopped() {
				limListener.Restart()
			}
		},
	}

	addr := s.listener.Addr().String()
	host, port, err := net.SplitHostPort(addr)
	if err != nil {
		panic("should not happen")
	}
	mimic.Host = host
	mimic.Port = port
	if chListenOn != nil {
		*chListenOn <- addr
	}

	return s.httpServer.Serve(s.listener)
}
<commit_msg>Redundant but clear<commit_after>package main

import (
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"os"
	"sync/atomic"
	"time"

	"github.com/gorilla/context"

	"github.com/getlantern/measured"

	"github.com/getlantern/http-proxy-extensions/devicefilter"
	"github.com/getlantern/http-proxy-extensions/mimic"
	"github.com/getlantern/http-proxy-extensions/profilter"
	"github.com/getlantern/http-proxy-extensions/tokenfilter"
	"github.com/getlantern/http-proxy/commonfilter"
	"github.com/getlantern/http-proxy/forward"
	"github.com/getlantern/http-proxy/httpconnect"
	"github.com/getlantern/http-proxy/utils"
)

type Server struct {
	firstHandler http.Handler
	httpServer   http.Server
	tls          bool

	listener net.Listener

	maxConns uint64
	numConns uint64

	idleTimeout time.Duration
}

func NewServer(token string, maxConns uint64, idleTimeout time.Duration, enableFilters bool, logLevel utils.LogLevel) *Server {
	stdWriter := io.Writer(os.Stdout)

	if maxConns == 0 {
		maxConns = math.MaxInt64
	}

	// The following middleware architecture can be seen as a chain of
	// filters that is run from last to first.
	// Don't forget to check Oxy and Gorilla's handlers for middleware.

	// Handles Direct Proxying
	forwardHandler, _ := forward.New(
		nil,
		forward.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		forward.IdleTimeoutSetter(idleTimeout),
	)

	// Handles HTTP CONNECT
	connectHandler, _ := httpconnect.New(
		forwardHandler,
		httpconnect.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		httpconnect.IdleTimeoutSetter(idleTimeout),
	)

	// Catches any request before reaching the CONNECT middleware or
	// the forwarder
	commonFilter, _ := commonfilter.New(
		connectHandler,
		commonfilter.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
	)

	var firstHandler http.Handler
	if !enableFilters {
		firstHandler = commonFilter
	} else {
		// Identifies Lantern Pro users (currently NOOP)
		lanternPro, _ := profilter.New(
			commonFilter,
			profilter.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		)
		// Returns a 404 to requests without the proper token.  Removes the
		// header before continuing.
		tokenFilter, _ := tokenfilter.New(
			lanternPro,
			tokenfilter.TokenSetter(token),
			tokenfilter.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		)
		// Extracts the user ID and attaches the matching client to the request
		// context.  Returns a 404 to requests without the UID.  Removes the
		// header before continuing.
		deviceFilter, _ := devicefilter.New(
			tokenFilter,
			devicefilter.Logger(utils.NewTimeLogger(&stdWriter, logLevel)),
		)
		firstHandler = deviceFilter
	}

	server := &Server{
		firstHandler: firstHandler,
		maxConns:     maxConns,
		numConns:     0,
		idleTimeout:  idleTimeout,
	}
	return server
}

func (s *Server) ServeHTTP(addr string, chListenOn *chan string) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	s.tls = false
	fmt.Printf("Listen http on %s\n", addr)
	return s.doServe(listener, chListenOn)
}

func (s *Server) ServeHTTPS(addr, keyfile, certfile string, chListenOn *chan string) error {
	listener, err := listenTLS(addr, keyfile, certfile)
	if err != nil {
		return err
	}
	s.tls = true
	fmt.Printf("Listen http on %s\n", addr)
	return s.doServe(listener, chListenOn)
}

func (s *Server) doServe(listener net.Listener, chListenOn *chan string) error {
	// A dirty trick to associate a connection with the http.Request it
	// contains. In "net/http/server.go", handler will be called
	// immediately after ConnState changed to StateActive, so it's safe to
	// loop through all elements in a channel to find a match remote addr.
	q := make(chan net.Conn, 10)

	proxy := http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			for c := range q {
				if c.RemoteAddr().String() == req.RemoteAddr {
					context.Set(req, "conn", c)
					break
				} else {
					q <- c
				}
			}
			s.firstHandler.ServeHTTP(w, req)
		})

	limListener := newLimitedListener(listener, &s.numConns, s.idleTimeout)
	if *enableReports {
		mListener := measured.Listener(limListener, 30*time.Second)
		s.listener = mListener
	} else {
		s.listener = limListener
	}

	s.httpServer = http.Server{Handler: proxy,
		ConnState: func(c net.Conn, state http.ConnState) {
			if state == http.StateActive {
				select {
				case q <- c:
				default:
					fmt.Print("Oops! the connection queue is full!\n")
				}
			}

			if atomic.LoadUint64(&s.numConns) >= s.maxConns {
				limListener.Stop()
			} else if limListener.IsStopped() {
				limListener.Restart()
			}
		},
	}

	addr := s.listener.Addr().String()
	host, port, err := net.SplitHostPort(addr)
	if err != nil {
		panic("should not happen")
	}
	mimic.Host = host
	mimic.Port = port
	if chListenOn != nil {
		*chListenOn <- addr
	}

	return s.httpServer.Serve(s.listener)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/kr/pty"
	"io"
	"log"
	"net"
	"net/rpc"
	"os"
	"os/exec"
	"sync"
	"time"
	"unicode/utf8"
)

type Console struct {
	file    *os.File
	command *exec.Cmd
}

type LockingConsoles struct {
	mutex sync.RWMutex
	byId  map[int64]Console
}

var consoles *LockingConsoles

func (c *LockingConsoles) deleteConsole(id int64) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	delete(c.byId, id)
}

func (c *LockingConsoles) addConsole(id int64, console Console) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	c.byId[id] = console
}

type ConsoleChunk struct {
	Id   int64
	Data []byte
}

var readChannel chan ConsoleChunk

func consoleReadLoop(output *os.File, id int64) {
	for {
		b := make([]byte, 1024)
		_, err := output.Read(b)
		if err == io.EOF {
			return
		}
		if err != nil {
			time.Sleep(100 * time.Millisecond)
			continue
		}
		out := fixUTF(string(b))
		readChannel <- ConsoleChunk{
			Id:   id,
			Data: []byte(out),
		}
	}
}

func fixUTF(s string) string {
	if !utf8.ValidString(s) {
		v := make([]rune, 0, len(s))
		for i, r := range s {
			if r == utf8.RuneError {
				_, size := utf8.DecodeRuneInString(s[i:])
				if size == 1 {
					continue
				}
			}
			v = append(v, r)
		}
		s = string(v)
	}
	return s
}

func consoleWriter(r io.Reader) {
	buffer := bufio.NewReader(r)
	for {
		str, err := buffer.ReadString('\n')
		if err != nil {
			continue
		}
		var chunk ConsoleChunk
		json.Unmarshal([]byte(str), &chunk)
		consoles.mutex.RLock()
		consoles.byId[chunk.Id].file.Write(chunk.Data)
		consoles.mutex.RUnlock()
	}
}

func consoleReader(c net.Conn) {
	for chunk := range readChannel {
		str, _ := json.Marshal(chunk)
		output := string(str) + "\n"
		_, err := c.Write([]byte(output))
		if err != nil {
			log.Println("Write: " + err.Error())
		}
	}
}

//console socket
func consoleListen() {
	l, err := net.Listen("unix", "@/tmp/vzconsole.sock")
	if err != nil {
		fmt.Println("listen error", err.Error())
		return
	}

	for {
		fd, err := l.Accept()
		if err != nil {
			fmt.Println("accept error", err.Error())
			return
		}
		go consoleReader(fd)
		go consoleWriter(fd)

	}
}

//rpc socket
func main() {
	log.Println("Started VZControl")
	consoles = &LockingConsoles{
		byId: make(map[int64]Console),
	}
	readChannel = make(chan ConsoleChunk)
	go consoleListen()
	vz := new(VZControl)
	rpc.Register(vz)
	listener, e := net.Listen("unix", "@/tmp/vzcontrol.sock")
	if e != nil {
		log.Fatal("listen error:", e)
	}

	for {
		if conn, err := listener.Accept(); err != nil {
			log.Fatal("accept error: " + err.Error())
		} else {
			log.Printf("new connection established\n")
			go rpc.ServeConn(conn)
		}
	}
}

type VZControl struct{}

func (vz *VZControl) ContainerCreate(cid int64, reply *int64) error {
	output, err := createContainer(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Create Error: %s\n Output:%s", err.Error(), output))
	}
	output, err = setupMount(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Mount Error: %s\n Output:%s", err.Error(), output))
	}
	output, err = startContainer(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Start Error: %s\n Output:%s", err.Error(), output))
	}
	reply = &cid
	return nil
}

func (vz *VZControl) ConsoleStart(cid int64, reply *int64) error {
	err := startConsole(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Console Start Error: %s", err.Error()))
	}
	reply = &cid
	return nil
}

func (vz *VZControl) ConsoleKill(cid int64, reply *int64) error {
	err := killConsole(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Console Kill Error: %s", err.Error()))
	}
	reply = &cid
	return nil
}

func (vz *VZControl) NetworkCreate(networkid int64, reply *int64) error {
	output, err := addBridge(networkid)
	if err != nil {
		return errors.New(fmt.Sprintf("Create Network Error: %s\n Output:%s", err.Error(), output))
	}
	reply = &networkid
	return nil
}

type NetworkAddArgs struct {
	Id, NetworkId int64
}

func (vz *VZControl) NetworkAdd(args *NetworkAddArgs, reply *int64) error {
	cid := args.Id
	networkid := args.NetworkId
	output, err := addInterface(cid, networkid)
	if err != nil {
		return errors.New(fmt.Sprintf("Interface Add Error: %s\n Output:%s", err.Error(), output))
	}
	output, err = connectBridge(cid, networkid)
	if err != nil {
		return errors.New(fmt.Sprintf("Bridge Connect Error: %s\n Output:%s", err.Error(), output))
	}
	reply = &cid
	return nil
}

func (vz *VZControl) Reset(someid int64, reply *int64) error {
	output, err := resetSystem()
	if err != nil {
		return errors.New(fmt.Sprintf("Reset Error: %s\n Output:%s", err.Error(), output))
	}
	consoles.mutex.RLock()
	consolesCopy := consoles.byId
	consoles.mutex.RUnlock()

	for cid, _ := range consolesCopy {
		killConsole(cid)
	}
	return nil
}

func resetSystem() (string, error) {
	command := exec.Command("./reset.sh")
	output, err := command.CombinedOutput()
	return string(output), err
}

func createContainer(id int64) (string, error) {
	command := exec.Command("vzctl", "create", fmt.Sprintf("%d", id), "--config", "ginux")
	output, err := command.CombinedOutput()
	return string(output), err
}

func setupMount(id int64) (string, error) {
	command := exec.Command("cp", "/etc/vz/conf/ginux.mount", fmt.Sprintf("/etc/vz/conf/%d.mount", id))
	output, err := command.CombinedOutput()
	return string(output), err
}

func startContainer(id int64) (string, error) {
	command := exec.Command("vzctl", "start", fmt.Sprintf("%d", id))
	output, err := command.CombinedOutput()
	return string(output), err
}

func addInterface(id int64, networkid int64) (string, error) {
	command := exec.Command("./addeth.sh", fmt.Sprintf("%d", id), fmt.Sprintf("%d", networkid))
	output, err := command.CombinedOutput()
	return string(output), err
}

func addBridge(networkid int64) (string, error) {
	command := exec.Command("./addbr.sh", fmt.Sprintf("%d", networkid))
	output, err := command.CombinedOutput()
	return string(output), err
}

func connectBridge(id int64, networkid int64) (string, error) {
	command := exec.Command("brctl", "addif", fmt.Sprintf("vzbr%d", networkid), fmt.Sprintf("veth%d.%d", id, networkid))
	output, err := command.CombinedOutput()
	return string(output), err
}

func startConsole(id int64) error {
	consoles.mutex.RLock()
	_, exists := consoles.byId[id]
	consoles.mutex.RUnlock()
	if exists {
		return errors.New(fmt.Sprintf("Console %d already is open", id))
	}
	cmd := exec.Command("vzctl", "console", fmt.Sprintf("%d", id))
	f, err := pty.Start(cmd)
	if err != nil {
		return err
	}
	con := Console{
		file:    f,
		command: cmd,
	}
	consoles.addConsole(id, con)
	go consoleReadLoop(f, id)
	return nil
}

func killConsole(id int64) error {
	consoles.mutex.RLock()
	console, ok := consoles.byId[id]
	consoles.mutex.RUnlock()
	if ok {
		if console.command != nil {
			err_kill := console.command.Process.Kill()
			if err_kill != nil {
				return err_kill
			}
			_, err_wait := console.command.Process.Wait()
			if err_wait != nil {
				return err_wait
			}
		}
	}
	consoles.deleteConsole(id)
	return nil
}
<commit_msg>Removed mount setup<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/kr/pty"
	"io"
	"log"
	"net"
	"net/rpc"
	"os"
	"os/exec"
	"sync"
	"time"
	"unicode/utf8"
)

type Console struct {
	file    *os.File
	command *exec.Cmd
}

type LockingConsoles struct {
	mutex sync.RWMutex
	byId  map[int64]Console
}

var consoles *LockingConsoles

func (c *LockingConsoles) deleteConsole(id int64) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	delete(c.byId, id)
}

func (c *LockingConsoles) addConsole(id int64, console Console) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	c.byId[id] = console
}

type ConsoleChunk struct {
	Id   int64
	Data []byte
}

var readChannel chan ConsoleChunk

func consoleReadLoop(output *os.File, id int64) {
	for {
		b := make([]byte, 1024)
		_, err := output.Read(b)
		if err == io.EOF {
			return
		}
		if err != nil {
			time.Sleep(100 * time.Millisecond)
			continue
		}
		out := fixUTF(string(b))
		readChannel <- ConsoleChunk{
			Id:   id,
			Data: []byte(out),
		}
	}
}

func fixUTF(s string) string {
	if !utf8.ValidString(s) {
		v := make([]rune, 0, len(s))
		for i, r := range s {
			if r == utf8.RuneError {
				_, size := utf8.DecodeRuneInString(s[i:])
				if size == 1 {
					continue
				}
			}
			v = append(v, r)
		}
		s = string(v)
	}
	return s
}

func consoleWriter(r io.Reader) {
	buffer := bufio.NewReader(r)
	for {
		str, err := buffer.ReadString('\n')
		if err != nil {
			continue
		}
		var chunk ConsoleChunk
		json.Unmarshal([]byte(str), &chunk)
		consoles.mutex.RLock()
		consoles.byId[chunk.Id].file.Write(chunk.Data)
		consoles.mutex.RUnlock()
	}
}

func consoleReader(c net.Conn) {
	for chunk := range readChannel {
		str, _ := json.Marshal(chunk)
		output := string(str) + "\n"
		_, err := c.Write([]byte(output))
		if err != nil {
			log.Println("Write: " + err.Error())
		}
	}
}

//console socket
func consoleListen() {
	l, err := net.Listen("unix", "@/tmp/vzconsole.sock")
	if err != nil {
		fmt.Println("listen error", err.Error())
		return
	}

	for {
		fd, err := l.Accept()
		if err != nil {
			fmt.Println("accept error", err.Error())
			return
		}
		go consoleReader(fd)
		go consoleWriter(fd)

	}
}

//rpc socket
func main() {
	log.Println("Started VZControl")
	consoles = &LockingConsoles{
		byId: make(map[int64]Console),
	}
	readChannel = make(chan ConsoleChunk)
	go consoleListen()
	vz := new(VZControl)
	rpc.Register(vz)
	listener, e := net.Listen("unix", "@/tmp/vzcontrol.sock")
	if e != nil {
		log.Fatal("listen error:", e)
	}

	for {
		if conn, err := listener.Accept(); err != nil {
			log.Fatal("accept error: " + err.Error())
		} else {
			log.Printf("new connection established\n")
			go rpc.ServeConn(conn)
		}
	}
}

type VZControl struct{}

func (vz *VZControl) ContainerCreate(cid int64, reply *int64) error {
	output, err := createContainer(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Create Error: %s\n Output:%s", err.Error(), output))
	}
	/*output, err = setupMount(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Mount Error: %s\n Output:%s", err.Error(), output))
	}
	output, err = startContainer(cid)*/
	if err != nil {
		return errors.New(fmt.Sprintf("Start Error: %s\n Output:%s", err.Error(), output))
	}
	reply = &cid
	return nil
}

func (vz *VZControl) ConsoleStart(cid int64, reply *int64) error {
	err := startConsole(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Console Start Error: %s", err.Error()))
	}
	reply = &cid
	return nil
}

func (vz *VZControl) ConsoleKill(cid int64, reply *int64) error {
	err := killConsole(cid)
	if err != nil {
		return errors.New(fmt.Sprintf("Console Kill Error: %s", err.Error()))
	}
	reply = &cid
	return nil
}

func (vz *VZControl) NetworkCreate(networkid int64, reply *int64) error {
	output, err := addBridge(networkid)
	if err != nil {
		return errors.New(fmt.Sprintf("Create Network Error: %s\n Output:%s", err.Error(), output))
	}
	reply = &networkid
	return nil
}

type NetworkAddArgs struct {
	Id, NetworkId int64
}

func (vz *VZControl) NetworkAdd(args *NetworkAddArgs, reply *int64) error {
	cid := args.Id
	networkid := args.NetworkId
	output, err := addInterface(cid, networkid)
	if err != nil {
		return errors.New(fmt.Sprintf("Interface Add Error: %s\n Output:%s", err.Error(), output))
	}
	output, err = connectBridge(cid, networkid)
	if err != nil {
		return errors.New(fmt.Sprintf("Bridge Connect Error: %s\n Output:%s", err.Error(), output))
	}
	reply = &cid
	return nil
}

func (vz *VZControl) Reset(someid int64, reply *int64) error {
	output, err := resetSystem()
	if err != nil {
		return errors.New(fmt.Sprintf("Reset Error: %s\n Output:%s", err.Error(), output))
	}
	consoles.mutex.RLock()
	consolesCopy := consoles.byId
	consoles.mutex.RUnlock()

	for cid, _ := range consolesCopy {
		killConsole(cid)
	}
	return nil
}

func resetSystem() (string, error) {
	command := exec.Command("./reset.sh")
	output, err := command.CombinedOutput()
	return string(output), err
}

func createContainer(id int64) (string, error) {
	command := exec.Command("vzctl", "create", fmt.Sprintf("%d", id), "--config", "ginux")
	output, err := command.CombinedOutput()
	return string(output), err
}

func setupMount(id int64) (string, error) {
	command := exec.Command("cp", "/etc/vz/conf/ginux.mount", fmt.Sprintf("/etc/vz/conf/%d.mount", id))
	output, err := command.CombinedOutput()
	return string(output), err
}

func startContainer(id int64) (string, error) {
	command := exec.Command("vzctl", "start", fmt.Sprintf("%d", id))
	output, err := command.CombinedOutput()
	return string(output), err
}

func addInterface(id int64, networkid int64) (string, error) {
	command := exec.Command("./addeth.sh", fmt.Sprintf("%d", id), fmt.Sprintf("%d", networkid))
	output, err := command.CombinedOutput()
	return string(output), err
}

func addBridge(networkid int64) (string, error) {
	command := exec.Command("./addbr.sh", fmt.Sprintf("%d", networkid))
	output, err := command.CombinedOutput()
	return string(output), err
}

func connectBridge(id int64, networkid int64) (string, error) {
	command := exec.Command("brctl", "addif", fmt.Sprintf("vzbr%d", networkid), fmt.Sprintf("veth%d.%d", id, networkid))
	output, err := command.CombinedOutput()
	return string(output), err
}

func startConsole(id int64) error {
	consoles.mutex.RLock()
	_, exists := consoles.byId[id]
	consoles.mutex.RUnlock()
	if exists {
		return errors.New(fmt.Sprintf("Console %d already is open", id))
	}
	cmd := exec.Command("vzctl", "console", fmt.Sprintf("%d", id))
	f, err := pty.Start(cmd)
	if err != nil {
		return err
	}
	con := Console{
		file:    f,
		command: cmd,
	}
	consoles.addConsole(id, con)
	go consoleReadLoop(f, id)
	return nil
}

func killConsole(id int64) error {
	consoles.mutex.RLock()
	console, ok := consoles.byId[id]
	consoles.mutex.RUnlock()
	if ok {
		if console.command != nil {
			err_kill := console.command.Process.Kill()
			if err_kill != nil {
				return err_kill
			}
			_, err_wait := console.command.Process.Wait()
			if err_wait != nil {
				return err_wait
			}
		}
	}
	consoles.deleteConsole(id)
	return nil
}
<|endoftext|>


<commit_before>package sshark

type JobId uint32
type MappedPort uint32

type StreamOutput struct {
	Name string
	Data string

	Finished   bool
	ExitStatus uint32
}

type JobInfo struct {
	ExitStatus uint32
}

type Container interface {
	Destroy() error
	Run(command string) (*JobInfo, error)
	NetIn() (MappedPort, error)
}
<commit_msg>remove types no longer used by Container interface<commit_after>package sshark

type MappedPort uint32

type JobInfo struct {
	ExitStatus uint32
}

type Container interface {
	Destroy() error
	Run(command string) (*JobInfo, error)
	NetIn() (MappedPort, error)
}
<|endoftext|>


<commit_before>package garden

import (
	"io"
	"time"
)

//go:generate counterfeiter . Container

type Container interface {
	Handle() string

	// Stop stops a container.
	//
	// If kill is false, garden stops a container by sending the processes running inside it the SIGTERM signal.
	// It then waits for the processes to terminate before returning a response.
	// If one or more processes do not terminate within 10 seconds,
	// garden sends these processes the SIGKILL signal, killing them ungracefully.
	//
	// If kill is true, garden stops a container by sending the processing running inside it a SIGKILL signal.
	//
	// It is possible to copy files in to and out of a stopped container.
	// It is only when a container is destroyed that its filesystem is cleaned up.
	//
	// Errors:
	// * None.
	Stop(kill bool) error

	// Returns information about a container.
	Info() (ContainerInfo, error)

	// StreamIn streams data into a file in a container.
	//
	// Errors:
	// *  TODO.
	StreamIn(spec StreamInSpec) error

	// StreamOut streams a file out of a container.
	//
	// Errors:
	// * TODO.
	StreamOut(spec StreamOutSpec) (io.ReadCloser, error)

	// Returns the current bandwidth limits set for the container.
	CurrentBandwidthLimits() (BandwidthLimits, error)

	// Returns the current CPU limts set for the container.
	CurrentCPULimits() (CPULimits, error)

	// Returns the current disk limts set for the container.
	CurrentDiskLimits() (DiskLimits, error)

	// Returns the current memory limts set for the container.
	CurrentMemoryLimits() (MemoryLimits, error)

	// Map a port on the host to a port in the container so that traffic to the
	// host port is forwarded to the container port. This is deprecated in
	// favour of passing NetIn configuration in the ContainerSpec at creation
	// time.
	//
	// If a host port is not given, a port will be acquired from the server's port
	// pool.
	//
	// If a container port is not given, the port will be the same as the
	// container port.
	//
	// The resulting host and container ports are returned in that order.
	//
	// Errors:
	// * When no port can be acquired from the server's port pool.
	NetIn(hostPort, containerPort uint32) (uint32, uint32, error)

	// Whitelist outbound network traffic. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// If the configuration directive deny_networks is not used,
	// all networks are already whitelisted and this command is effectively a no-op.
	//
	// Later NetOut calls take precedence over earlier calls, which is
	// significant only in relation to logging.
	//
	// Errors:
	// * An error is returned if the NetOut call fails.
	NetOut(netOutRule NetOutRule) error

	// A Bulk call for NetOut. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// Errors:
	// * An error is returned if any of the NetOut calls fail.
	BulkNetOut(netOutRules []NetOutRule) error

	// Run a script inside a container.
	//
	// The root user will be mapped to a non-root UID in the host unless the container (not this process) was created with 'privileged' true.
	//
	// Errors:
	// * TODO.
	Run(ProcessSpec, ProcessIO) (Process, error)

	// Attach starts streaming the output back to the client from a specified process.
	//
	// Errors:
	// * processID does not refer to a running process.
	Attach(processID string, io ProcessIO) (Process, error)

	// Metrics returns the current set of metrics for a container
	Metrics() (Metrics, error)

	// Sets the grace time.
	SetGraceTime(graceTime time.Duration) error

	// Properties returns the current set of properties
	Properties() (Properties, error)

	// Property returns the value of the property with the specified name.
	//
	// Errors:
	// * When the property does not exist on the container.
	Property(name string) (string, error)

	// Set a named property on a container to a specified value.
	//
	// Errors:
	// * None.
	SetProperty(name string, value string) error

	// Remove a property with the specified name from a container.
	//
	// Errors:
	// * None.
	RemoveProperty(name string) error
}

// ProcessSpec contains parameters for running a script inside a container.
type ProcessSpec struct {
	// ID for the process. If empty, an ID will be generated.
	ID string `json:"id,omitempty"`

	// Path to command to execute.
	Path string `json:"path,omitempty"`

	// Arguments to pass to command.
	Args []string `json:"args,omitempty"`

	// Environment variables.
	Env []string `json:"env,omitempty"`

	// Working directory (default: home directory).
	Dir string `json:"dir,omitempty"`

	// The name of a user in the container to run the process as.
	User string `json:"user,omitempty"`

	// Resource limits
	Limits ResourceLimits `json:"rlimits,omitempty"`

	// Execute with a TTY for stdio.
	TTY *TTYSpec `json:"tty,omitempty"`

	// Execute process in own root filesystem, different from the other processes
	// in the container.
	Image ImageRef `json:"image,omitempty"`
}

type TTYSpec struct {
	WindowSize *WindowSize `json:"window_size,omitempty"`
}

type WindowSize struct {
	Columns int `json:"columns,omitempty"`
	Rows    int `json:"rows,omitempty"`
}

type ProcessIO struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

//go:generate counterfeiter . Process

type Process interface {
	ID() string
	Wait() (int, error)
	SetTTY(TTYSpec) error
	Signal(Signal) error
}

type Signal int

const (
	SignalTerminate Signal = iota
	SignalKill
)

type PortMapping struct {
	HostPort      uint32
	ContainerPort uint32
}

type StreamInSpec struct {
	Path      string
	User      string
	TarStream io.Reader
}

type StreamOutSpec struct {
	Path string
	User string
}

// ContainerInfo holds information about a container.
type ContainerInfo struct {
	State         string        // Either "active" or "stopped".
	Events        []string      // List of events that occurred for the container. It currently includes only "oom" (Out Of Memory) event if it occurred.
	HostIP        string        // The IP address of the gateway which controls the host side of the container's virtual ethernet pair.
	ContainerIP   string        // The IP address of the container side of the container's virtual ethernet pair.
	ExternalIP    string        //
	ContainerPath string        // The path to the directory holding the container's files (both its control scripts and filesystem).
	ProcessIDs    []string      // List of running processes.
	Properties    Properties    // List of properties defined for the container.
	MappedPorts   []PortMapping //
}

type ContainerInfoEntry struct {
	Info ContainerInfo
	Err  *Error
}

type Metrics struct {
	MemoryStat  ContainerMemoryStat
	CPUStat     ContainerCPUStat
	DiskStat    ContainerDiskStat
	NetworkStat ContainerNetworkStat
}

type ContainerMetricsEntry struct {
	Metrics Metrics
	Err     *Error
}

type ContainerMemoryStat struct {
	ActiveAnon              uint64 `json:"active_anon"`
	ActiveFile              uint64 `json:"active_file"`
	Cache                   uint64 `json:"cache"`
	HierarchicalMemoryLimit uint64 `json:"hierarchical_memory_limit"`
	InactiveAnon            uint64 `json:"inactive_anon"`
	InactiveFile            uint64 `json:"inactive_file"`
	MappedFile              uint64 `json:"mapped_file"`
	Pgfault                 uint64 `json:"pgfault"`
	Pgmajfault              uint64 `json:"pgmajfault"`
	Pgpgin                  uint64 `json:"pgpgin"`
	Pgpgout                 uint64 `json:"pgpgout"`
	Rss                     uint64 `json:"rss"`
	TotalActiveAnon         uint64 `json:"total_active_anon"`
	TotalActiveFile         uint64 `json:"total_active_file"`
	TotalCache              uint64 `json:"total_cache"`
	TotalInactiveAnon       uint64 `json:"total_inactive_anon"`
	TotalInactiveFile       uint64 `json:"total_inactive_file"`
	TotalMappedFile         uint64 `json:"total_mapped_file"`
	TotalPgfault            uint64 `json:"total_pgfault"`
	TotalPgmajfault         uint64 `json:"total_pgmajfault"`
	TotalPgpgin             uint64 `json:"total_pgpgin"`
	TotalPgpgout            uint64 `json:"total_pgpgout"`
	TotalRss                uint64 `json:"total_rss"`
	TotalUnevictable        uint64 `json:"total_unevictable"`
	Unevictable             uint64 `json:"unevictable"`
	Swap                    uint64 `json:"swap"`
	HierarchicalMemswLimit  uint64 `json:"hierarchical_memsw_limit"`
	TotalSwap               uint64 `json:"total_swap"`
	// A memory usage total which reports memory usage in the same way that limits are enforced.
	// This value includes memory consumed by nested containers.
	TotalUsageTowardLimit uint64
}

type ContainerCPUStat struct {
	Usage  uint64
	User   uint64
	System uint64
}

type ContainerDiskStat struct {
	TotalBytesUsed      uint64
	TotalInodesUsed     uint64
	ExclusiveBytesUsed  uint64
	ExclusiveInodesUsed uint64
}

type ContainerBandwidthStat struct {
	InRate   uint64
	InBurst  uint64
	OutRate  uint64
	OutBurst uint64
}

type ContainerNetworkStat struct {
	RxBytes uint64
	TxBytes uint64
}

type BandwidthLimits struct {
	RateInBytesPerSecond      uint64 `json:"rate,omitempty"`
	BurstRateInBytesPerSecond uint64 `json:"burst,omitempty"`
}

type DiskLimits struct {
	InodeSoft uint64 `json:"inode_soft,omitempty"`
	InodeHard uint64 `json:"inode_hard,omitempty"`

	ByteSoft uint64 `json:"byte_soft,omitempty"`
	ByteHard uint64 `json:"byte_hard,omitempty"`

	Scope DiskLimitScope `json:"scope,omitempty"`
}

type MemoryLimits struct {
	//	Memory usage limit in bytes.
	LimitInBytes uint64 `json:"limit_in_bytes,omitempty"`
}

type CPULimits struct {
	LimitInShares uint64 `json:"limit_in_shares,omitempty"`
}

type PidLimits struct {
	// Limits the number of pids a container may create before new forks or clones are disallowed to processes in the container.
	// Note: this may only be enforced when a process attempts to fork, so it does not guarantee that a new container.Run(ProcessSpec)
	// will not succeed even if the limit has been exceeded, but the process will not be able to spawn further processes or threads.
	Max uint64 `json:"max,omitempty"`
}

// Resource limits.
//
// Please refer to the manual page of getrlimit for a description of the individual fields:
// http://www.kernel.org/doc/man-pages/online/pages/man2/getrlimit.2.html
type ResourceLimits struct {
	As         *uint64 `json:"as,omitempty"`
	Core       *uint64 `json:"core,omitempty"`
	Cpu        *uint64 `json:"cpu,omitempty"`
	Data       *uint64 `json:"data,omitempty"`
	Fsize      *uint64 `json:"fsize,omitempty"`
	Locks      *uint64 `json:"locks,omitempty"`
	Memlock    *uint64 `json:"memlock,omitempty"`
	Msgqueue   *uint64 `json:"msgqueue,omitempty"`
	Nice       *uint64 `json:"nice,omitempty"`
	Nofile     *uint64 `json:"nofile,omitempty"`
	Nproc      *uint64 `json:"nproc,omitempty"`
	Rss        *uint64 `json:"rss,omitempty"`
	Rtprio     *uint64 `json:"rtprio,omitempty"`
	Sigpending *uint64 `json:"sigpending,omitempty"`
	Stack      *uint64 `json:"stack,omitempty"`
}

type DiskLimitScope uint8

const DiskLimitScopeTotal DiskLimitScope = 0
const DiskLimitScopeExclusive DiskLimitScope = 1
<commit_msg>Document uid:gid format for ProcessSpec.User<commit_after>package garden

import (
	"io"
	"time"
)

//go:generate counterfeiter . Container

type Container interface {
	Handle() string

	// Stop stops a container.
	//
	// If kill is false, garden stops a container by sending the processes running inside it the SIGTERM signal.
	// It then waits for the processes to terminate before returning a response.
	// If one or more processes do not terminate within 10 seconds,
	// garden sends these processes the SIGKILL signal, killing them ungracefully.
	//
	// If kill is true, garden stops a container by sending the processing running inside it a SIGKILL signal.
	//
	// It is possible to copy files in to and out of a stopped container.
	// It is only when a container is destroyed that its filesystem is cleaned up.
	//
	// Errors:
	// * None.
	Stop(kill bool) error

	// Returns information about a container.
	Info() (ContainerInfo, error)

	// StreamIn streams data into a file in a container.
	//
	// Errors:
	// *  TODO.
	StreamIn(spec StreamInSpec) error

	// StreamOut streams a file out of a container.
	//
	// Errors:
	// * TODO.
	StreamOut(spec StreamOutSpec) (io.ReadCloser, error)

	// Returns the current bandwidth limits set for the container.
	CurrentBandwidthLimits() (BandwidthLimits, error)

	// Returns the current CPU limts set for the container.
	CurrentCPULimits() (CPULimits, error)

	// Returns the current disk limts set for the container.
	CurrentDiskLimits() (DiskLimits, error)

	// Returns the current memory limts set for the container.
	CurrentMemoryLimits() (MemoryLimits, error)

	// Map a port on the host to a port in the container so that traffic to the
	// host port is forwarded to the container port. This is deprecated in
	// favour of passing NetIn configuration in the ContainerSpec at creation
	// time.
	//
	// If a host port is not given, a port will be acquired from the server's port
	// pool.
	//
	// If a container port is not given, the port will be the same as the
	// container port.
	//
	// The resulting host and container ports are returned in that order.
	//
	// Errors:
	// * When no port can be acquired from the server's port pool.
	NetIn(hostPort, containerPort uint32) (uint32, uint32, error)

	// Whitelist outbound network traffic. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// If the configuration directive deny_networks is not used,
	// all networks are already whitelisted and this command is effectively a no-op.
	//
	// Later NetOut calls take precedence over earlier calls, which is
	// significant only in relation to logging.
	//
	// Errors:
	// * An error is returned if the NetOut call fails.
	NetOut(netOutRule NetOutRule) error

	// A Bulk call for NetOut. This is deprecated in favour of passing
	// NetOut configuration in the ContainerSpec at creation time.
	//
	// Errors:
	// * An error is returned if any of the NetOut calls fail.
	BulkNetOut(netOutRules []NetOutRule) error

	// Run a script inside a container.
	//
	// The root user will be mapped to a non-root UID in the host unless the container (not this process) was created with 'privileged' true.
	//
	// Errors:
	// * TODO.
	Run(ProcessSpec, ProcessIO) (Process, error)

	// Attach starts streaming the output back to the client from a specified process.
	//
	// Errors:
	// * processID does not refer to a running process.
	Attach(processID string, io ProcessIO) (Process, error)

	// Metrics returns the current set of metrics for a container
	Metrics() (Metrics, error)

	// Sets the grace time.
	SetGraceTime(graceTime time.Duration) error

	// Properties returns the current set of properties
	Properties() (Properties, error)

	// Property returns the value of the property with the specified name.
	//
	// Errors:
	// * When the property does not exist on the container.
	Property(name string) (string, error)

	// Set a named property on a container to a specified value.
	//
	// Errors:
	// * None.
	SetProperty(name string, value string) error

	// Remove a property with the specified name from a container.
	//
	// Errors:
	// * None.
	RemoveProperty(name string) error
}

// ProcessSpec contains parameters for running a script inside a container.
type ProcessSpec struct {
	// ID for the process. If empty, an ID will be generated.
	ID string `json:"id,omitempty"`

	// Path to command to execute.
	Path string `json:"path,omitempty"`

	// Arguments to pass to command.
	Args []string `json:"args,omitempty"`

	// Environment variables.
	Env []string `json:"env,omitempty"`

	// Working directory (default: home directory).
	Dir string `json:"dir,omitempty"`

	// The name of a user in the container to run the process as.
	// This must either be a username, or uid:gid.
	User string `json:"user,omitempty"`

	// Resource limits
	Limits ResourceLimits `json:"rlimits,omitempty"`

	// Execute with a TTY for stdio.
	TTY *TTYSpec `json:"tty,omitempty"`

	// Execute process in own root filesystem, different from the other processes
	// in the container.
	Image ImageRef `json:"image,omitempty"`
}

type TTYSpec struct {
	WindowSize *WindowSize `json:"window_size,omitempty"`
}

type WindowSize struct {
	Columns int `json:"columns,omitempty"`
	Rows    int `json:"rows,omitempty"`
}

type ProcessIO struct {
	Stdin  io.Reader
	Stdout io.Writer
	Stderr io.Writer
}

//go:generate counterfeiter . Process

type Process interface {
	ID() string
	Wait() (int, error)
	SetTTY(TTYSpec) error
	Signal(Signal) error
}

type Signal int

const (
	SignalTerminate Signal = iota
	SignalKill
)

type PortMapping struct {
	HostPort      uint32
	ContainerPort uint32
}

type StreamInSpec struct {
	Path      string
	User      string
	TarStream io.Reader
}

type StreamOutSpec struct {
	Path string
	User string
}

// ContainerInfo holds information about a container.
type ContainerInfo struct {
	State         string        // Either "active" or "stopped".
	Events        []string      // List of events that occurred for the container. It currently includes only "oom" (Out Of Memory) event if it occurred.
	HostIP        string        // The IP address of the gateway which controls the host side of the container's virtual ethernet pair.
	ContainerIP   string        // The IP address of the container side of the container's virtual ethernet pair.
	ExternalIP    string        //
	ContainerPath string        // The path to the directory holding the container's files (both its control scripts and filesystem).
	ProcessIDs    []string      // List of running processes.
	Properties    Properties    // List of properties defined for the container.
	MappedPorts   []PortMapping //
}

type ContainerInfoEntry struct {
	Info ContainerInfo
	Err  *Error
}

type Metrics struct {
	MemoryStat  ContainerMemoryStat
	CPUStat     ContainerCPUStat
	DiskStat    ContainerDiskStat
	NetworkStat ContainerNetworkStat
}

type ContainerMetricsEntry struct {
	Metrics Metrics
	Err     *Error
}

type ContainerMemoryStat struct {
	ActiveAnon              uint64 `json:"active_anon"`
	ActiveFile              uint64 `json:"active_file"`
	Cache                   uint64 `json:"cache"`
	HierarchicalMemoryLimit uint64 `json:"hierarchical_memory_limit"`
	InactiveAnon            uint64 `json:"inactive_anon"`
	InactiveFile            uint64 `json:"inactive_file"`
	MappedFile              uint64 `json:"mapped_file"`
	Pgfault                 uint64 `json:"pgfault"`
	Pgmajfault              uint64 `json:"pgmajfault"`
	Pgpgin                  uint64 `json:"pgpgin"`
	Pgpgout                 uint64 `json:"pgpgout"`
	Rss                     uint64 `json:"rss"`
	TotalActiveAnon         uint64 `json:"total_active_anon"`
	TotalActiveFile         uint64 `json:"total_active_file"`
	TotalCache              uint64 `json:"total_cache"`
	TotalInactiveAnon       uint64 `json:"total_inactive_anon"`
	TotalInactiveFile       uint64 `json:"total_inactive_file"`
	TotalMappedFile         uint64 `json:"total_mapped_file"`
	TotalPgfault            uint64 `json:"total_pgfault"`
	TotalPgmajfault         uint64 `json:"total_pgmajfault"`
	TotalPgpgin             uint64 `json:"total_pgpgin"`
	TotalPgpgout            uint64 `json:"total_pgpgout"`
	TotalRss                uint64 `json:"total_rss"`
	TotalUnevictable        uint64 `json:"total_unevictable"`
	Unevictable             uint64 `json:"unevictable"`
	Swap                    uint64 `json:"swap"`
	HierarchicalMemswLimit  uint64 `json:"hierarchical_memsw_limit"`
	TotalSwap               uint64 `json:"total_swap"`
	// A memory usage total which reports memory usage in the same way that limits are enforced.
	// This value includes memory consumed by nested containers.
	TotalUsageTowardLimit uint64
}

type ContainerCPUStat struct {
	Usage  uint64
	User   uint64
	System uint64
}

type ContainerDiskStat struct {
	TotalBytesUsed      uint64
	TotalInodesUsed     uint64
	ExclusiveBytesUsed  uint64
	ExclusiveInodesUsed uint64
}

type ContainerBandwidthStat struct {
	InRate   uint64
	InBurst  uint64
	OutRate  uint64
	OutBurst uint64
}

type ContainerNetworkStat struct {
	RxBytes uint64
	TxBytes uint64
}

type BandwidthLimits struct {
	RateInBytesPerSecond      uint64 `json:"rate,omitempty"`
	BurstRateInBytesPerSecond uint64 `json:"burst,omitempty"`
}

type DiskLimits struct {
	InodeSoft uint64 `json:"inode_soft,omitempty"`
	InodeHard uint64 `json:"inode_hard,omitempty"`

	ByteSoft uint64 `json:"byte_soft,omitempty"`
	ByteHard uint64 `json:"byte_hard,omitempty"`

	Scope DiskLimitScope `json:"scope,omitempty"`
}

type MemoryLimits struct {
	//	Memory usage limit in bytes.
	LimitInBytes uint64 `json:"limit_in_bytes,omitempty"`
}

type CPULimits struct {
	LimitInShares uint64 `json:"limit_in_shares,omitempty"`
}

type PidLimits struct {
	// Limits the number of pids a container may create before new forks or clones are disallowed to processes in the container.
	// Note: this may only be enforced when a process attempts to fork, so it does not guarantee that a new container.Run(ProcessSpec)
	// will not succeed even if the limit has been exceeded, but the process will not be able to spawn further processes or threads.
	Max uint64 `json:"max,omitempty"`
}

// Resource limits.
//
// Please refer to the manual page of getrlimit for a description of the individual fields:
// http://www.kernel.org/doc/man-pages/online/pages/man2/getrlimit.2.html
type ResourceLimits struct {
	As         *uint64 `json:"as,omitempty"`
	Core       *uint64 `json:"core,omitempty"`
	Cpu        *uint64 `json:"cpu,omitempty"`
	Data       *uint64 `json:"data,omitempty"`
	Fsize      *uint64 `json:"fsize,omitempty"`
	Locks      *uint64 `json:"locks,omitempty"`
	Memlock    *uint64 `json:"memlock,omitempty"`
	Msgqueue   *uint64 `json:"msgqueue,omitempty"`
	Nice       *uint64 `json:"nice,omitempty"`
	Nofile     *uint64 `json:"nofile,omitempty"`
	Nproc      *uint64 `json:"nproc,omitempty"`
	Rss        *uint64 `json:"rss,omitempty"`
	Rtprio     *uint64 `json:"rtprio,omitempty"`
	Sigpending *uint64 `json:"sigpending,omitempty"`
	Stack      *uint64 `json:"stack,omitempty"`
}

type DiskLimitScope uint8

const DiskLimitScopeTotal DiskLimitScope = 0
const DiskLimitScopeExclusive DiskLimitScope = 1
<|endoftext|>


<commit_before>package test

import (
	"io"
	"net"
	"testing"
	"time"

	"github.com/neptulon/client"
)

// ClientHelper is a client.Client wrapper for testing.
// All the functions are wrapped with proper test runner error logging.
type ClientHelper struct {
	client    *client.Client
	server    *ServerHelper // server that this connection will be made to
	testing   *testing.T
	cert, key []byte
}

// NewClientHelper creates a new client helper object.
// Takes target server as an argument to retrieve server certs, address, etc.
func NewClientHelper(t *testing.T, s *ServerHelper) *ClientHelper {
	if testing.Short() {
		t.Skip("Skipping integration test in short testing mode")
	}

	return &ClientHelper{testing: t, server: s}
}

// DialTLS initiates a TLS connection.
func (c *ClientHelper) DialTLS() *ClientHelper {
	// retry connect in case we're operating on a very slow machine
	for i := 0; i <= 5; i++ {
		conn, err := client.DialTLS(addr, c.server.IntCACert, c.cert, c.key, false) // no need for debug mode on conn as we have it on server conn already
		if err != nil {
			if operr, ok := err.(*net.OpError); ok && operr.Op == "dial" && operr.Err.Error() == "connection refused" {
				time.Sleep(time.Millisecond * 50)
				continue
			} else if i == 5 {
				c.testing.Fatalf("Cannot connect to server address %v after 5 retries, with error: %v", addr, err)
			}
			c.testing.Fatalf("Cannot connect to server address %v with error: %v", addr, err)
		}

		if i != 0 {
			c.testing.Logf("WARNING: it took %v retries to connect to the server, which might indicate code issues or slow machine.", i)
		}

		conn.SetReadDeadline(10)
		c.conn = conn
		return c
	}
}

// VerifyConnClosed verifies that the connection is in closed state.
// Verification is done via reading from the channel and checking that returned error is io.EOF.
func (c *ConnHelper) VerifyConnClosed() bool {
	_, _, _, err := c.conn.ReadMsg(nil, nil)
	if err != io.EOF {
		return false
	}

	return true
}

// Close closes a connection.
func (c *ConnHelper) Close() {
	if err := c.conn.Close(); err != nil {
		c.testing.Fatal("Failed to close connection:", err)
	}
}
<commit_msg>remove comment<commit_after>package test

import (
	"io"
	"net"
	"testing"
	"time"

	"github.com/neptulon/client"
)

// ClientHelper is a client.Client wrapper for testing.
// All the functions are wrapped with proper test runner error logging.
type ClientHelper struct {
	client    *client.Client
	server    *ServerHelper // server that this connection will be made to
	testing   *testing.T
	cert, key []byte
}

// NewClientHelper creates a new client helper object.
// Takes target server as an argument to retrieve server certs, address, etc.
func NewClientHelper(t *testing.T, s *ServerHelper) *ClientHelper {
	if testing.Short() {
		t.Skip("Skipping integration test in short testing mode")
	}

	return &ClientHelper{testing: t, server: s}
}

// DialTLS initiates a TLS connection.
func (c *ClientHelper) DialTLS() *ClientHelper {
	// retry connect in case we're operating on a very slow machine
	for i := 0; i <= 5; i++ {
		conn, err := client.DialTLS(addr, c.server.IntCACert, c.cert, c.key, false)
		if err != nil {
			if operr, ok := err.(*net.OpError); ok && operr.Op == "dial" && operr.Err.Error() == "connection refused" {
				time.Sleep(time.Millisecond * 50)
				continue
			} else if i == 5 {
				c.testing.Fatalf("Cannot connect to server address %v after 5 retries, with error: %v", addr, err)
			}
			c.testing.Fatalf("Cannot connect to server address %v with error: %v", addr, err)
		}

		if i != 0 {
			c.testing.Logf("WARNING: it took %v retries to connect to the server, which might indicate code issues or slow machine.", i)
		}

		conn.SetReadDeadline(10)
		c.conn = conn
		return c
	}
}

// VerifyConnClosed verifies that the connection is in closed state.
// Verification is done via reading from the channel and checking that returned error is io.EOF.
func (c *ConnHelper) VerifyConnClosed() bool {
	_, _, _, err := c.conn.ReadMsg(nil, nil)
	if err != io.EOF {
		return false
	}

	return true
}

// Close closes a connection.
func (c *ConnHelper) Close() {
	if err := c.conn.Close(); err != nil {
		c.testing.Fatal("Failed to close connection:", err)
	}
}
<|endoftext|>


<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2014 AT&T
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package impl

import (
	"errors"
	"os"
	"fmt"
	"time"
	"strings"

	"code.google.com/p/gogoprotobuf/proto"
	"github.com/golang/glog"

	"github.com/att-innovate/charmander-scheduler/communication"
	"github.com/att-innovate/charmander-scheduler/mesosproto"
	"github.com/att-innovate/charmander-scheduler/scheduler"
	"github.com/att-innovate/charmander-scheduler/upid"

	managerInterface "github.com/att-innovate/charmander-scheduler/manager"
)

var taskRegistry = NewTaskRegistry()
var nodeRegistry = NewNodeRegistry()

func New(
	scheduler *scheduler.Scheduler,
	frameworkName string,
	master string,
	listening string) (manager, error) {

	framework := &mesosproto.FrameworkInfo{
		User: proto.String(""),
		Name: proto.String(frameworkName),
	}

	newManager := manager{
		scheduler:     scheduler,
		frameworkInfo: framework,
		master:        master,
		listening:     listening,
	}

	return newManager, nil
}

type manager struct {
	scheduler     *scheduler.Scheduler
	frameworkInfo *mesosproto.FrameworkInfo
	masterUPID    *upid.UPID
	selfUPID      *upid.UPID
	frameworkId   string
	master        string
	listening     string
}

func (self *manager) Start() error {
	retryCounter := 10

	for ; retryCounter > 0; retryCounter-- {
		if communication.MesosMasterReachable(self.master) { break }
		time.Sleep(6 * time.Second)
	}

	if retryCounter == 0 {
		return errors.New("Mesos unreachable")
	}

	self.frameworkInfo.User = proto.String("root")

	// set default hostname
	if self.frameworkInfo.GetHostname() == "" {
		host, err := os.Hostname()
		if err != nil || host == "" {
			host = "unknown"
		}
		self.frameworkInfo.Hostname = proto.String(host)
	}

	if m, err := upid.Parse("master@" + self.master); err != nil {
		return err
	} else {
		self.masterUPID = m
	}

	self.selfUPID= &upid.UPID{
		ID: "scheduler",
		Host: self.GetListenerIP(),
		Port: fmt.Sprintf("%d", self.GetListenerPortForScheduler())}

	communication.InitRestHandler(self)

	self.announceFramework()

	return nil
}

func (self *manager) GetListenerIP() string {
	return self.listening
}

func (self *manager) GetListenerPortForScheduler() int {
	return 7070
}

func (self *manager) GetListenerPortForRESTApi() int {
	return 7075
}

func (self *manager) GetTaskRequests() []*managerInterface.Task {
	return taskRegistry.Tasks()
}

func (self *manager) ResourceRequirementsWouldMatch(offer *mesosproto.Offer, taskRequest *managerInterface.Task) bool {
	for _, resource := range offer.GetResources() {
		switch {
		case *resource.Name == "cpus":
			if resource.Scalar.GetValue() < float64(taskRequest.Cpus) { return false }
		case *resource.Name == "mem":
			if resource.Scalar.GetValue() < float64(taskRequest.Mem) { return false }
		}
	}
	for _, attribute := range offer.GetAttributes() {
		switch {
		case *attribute.Name == "nodetype":
			if len(taskRequest.NodeType) > 0 && attribute.Text.GetValue() != taskRequest.NodeType { return false }
		case *attribute.Name == "nodename":
			if len(taskRequest.NodeName) > 0 && attribute.Text.GetValue() != taskRequest.NodeName { return false }
		}
	}

	return true
}

func (self *manager) HandleFrameworkRegistered(frameworkId string) {
	self.frameworkId = frameworkId
	if self.scheduler != nil && self.scheduler.Registered != nil {
		self.scheduler.Registered(self, frameworkId)
	}
}

func (self *manager) HandleResourceOffered(offers []*mesosproto.Offer) {
	self.updateNodeRegistry(offers)

	offers= self.enforceSLAs(offers)

	if self.scheduler != nil && self.scheduler.ResourceOffers != nil {
		self.scheduler.ResourceOffers(self, offers)
	}
}

func (self *manager) updateNodeRegistry(offers []*mesosproto.Offer) {
	for _, offer := range offers {
		slaveID := offer.GetSlaveId().GetValue()
		if nodeRegistry.Exists(slaveID) {
			nodeRegistry.UpdateTimeOfLastOffer(slaveID)
			continue
		}

		node := &managerInterface.Node {
			ID: slaveID,
			Hostname: offer.GetHostname(),
			TimeOfLastOffer: time.Now().Unix(),
		}

		for _, attribute := range offer.GetAttributes() {
			if *attribute.Name == "nodename" {
				node.NodeName = attribute.Text.GetValue()
			} else if *attribute.Name == "nodetype" {
				node.NodeType = attribute.Text.GetValue()
			}
		}

		nodeRegistry.Register(slaveID, node)

		glog.Infoln("New Node Registered: ", node)
	}
}

func (self *manager) enforceSLAs(offers []*mesosproto.Offer) []*mesosproto.Offer {
	result := offers
	var taskRequests []*managerInterface.Task
	taskRequests = self.GetTaskRequests()

	for _, taskRequest := range taskRequests {
		if taskRequest.Running { continue }
		if len(taskRequest.Sla) == 0 {continue }

		if taskRequest.Sla == managerInterface.SLA_ONE_PER_NODE {
			for _, offer := range offers {
				if resolveNodeName(*offer) == taskRequest.NodeName {
					if self.ResourceRequirementsWouldMatch(offer, taskRequest) {
						self.AcceptOffer(offer.GetId(), offer.SlaveId, taskRequest)
						result = removeOfferFromList(result, offer)
					}
				}
			}
		}
	}

	return result
}

func resolveNodeName(offer mesosproto.Offer) string {
	for _, attribute := range offer.GetAttributes() {
		if *attribute.Name == "nodename" {
			return attribute.Text.GetValue()
		}
	}

	return "notfound"
}

func removeOfferFromList(offers []*mesosproto.Offer, offer *mesosproto.Offer) []*mesosproto.Offer {
	numberOfOffers := len(offers)
	if numberOfOffers == 1 { return []*mesosproto.Offer {}}

	result := make([]*mesosproto.Offer, numberOfOffers-1)
	i :=  0
	for _, value := range offers {
		if offer.GetId().GetValue() == value.GetId().GetValue() { continue }
		result[i] = value
		i++
	}

	return result
}

func (self *manager) HandleStatusMessage(statusMessage *mesosproto.StatusUpdateMessage) {
	glog.Infof("Status Update %v\n", statusMessage)
	status := statusMessage.GetUpdate().GetStatus()

	switch {
	case  *status.State == mesosproto.TaskState_TASK_RUNNING:
		task, _ := taskRegistry.Fetch(status.GetTaskId().GetValue())
		task.Running= true
		task.SlaveID= status.GetSlaveId().GetValue()
		glog.Infof("Running: %s on Slave: %s", task.InternalID, task.SlaveID)
	case  *status.State == mesosproto.TaskState_TASK_FAILED:
		taskRegistry.Delete(status.GetTaskId().GetValue())
		glog.Infoln("Task Failed: ", status.GetTaskId().GetValue())
	case  *status.State == mesosproto.TaskState_TASK_LOST:
		taskRegistry.Delete(status.GetTaskId().GetValue())
		glog.Infoln("Task Lost: ", status.GetTaskId().GetValue())
	case  *status.State == mesosproto.TaskState_TASK_FINISHED:
		taskRegistry.Delete(status.GetTaskId().GetValue())
		glog.Infoln("Task Finished: ", status.GetTaskId().GetValue())
	}

	self.acknowledgeStatusUpdate(statusMessage)
}

func (self *manager) HandleRunDockerImage(task *managerInterface.Task) {
	if task.Sla == managerInterface.SLA_ONE_PER_NODE {
		for _, node := range nodeRegistry.Nodes() {
			newTask := managerInterface.CopyTask(*task)
			newTask.NodeName= node.NodeName
			self.handleRunDockerImageImpl(&newTask)
		}
	} else {
		self.handleRunDockerImageImpl(task)
	}
}

func (self *manager) handleRunDockerImageImpl(task *managerInterface.Task) {
	id := fmt.Sprintf("%v-%v", strings.Replace(task.ID, " ", "", -1), time.Now().UnixNano())
	memory := float64(task.Mem)
	portResources := []*mesosproto.Value_Range{}
	cpus := float64(0.1)
	if task.Cpus > 0 {
		cpus= task.Cpus
	}

	dockerInfo := &mesosproto.ContainerInfo_DockerInfo {
		Image: &task.DockerImage,
	}
	containerInfo := &mesosproto.ContainerInfo {
		Type: mesosproto.ContainerInfo_DOCKER.Enum(),
		Docker: dockerInfo,
	}
	for _, volume := range task.Volumes {
		mode := mesosproto.Volume_RW
		if volume.Mode == "ro" {
			mode = mesosproto.Volume_RO
		}

		containerInfo.Volumes = append(containerInfo.Volumes, &mesosproto.Volume{
				ContainerPath: &volume.ContainerPath,
				HostPath:      &volume.HostPath,
				Mode:          &mode,
			})
	}
	for _, port := range task.Ports {
		dockerInfo.PortMappings = append(dockerInfo.PortMappings, &mesosproto.ContainerInfo_DockerInfo_PortMapping{
				ContainerPort: &port.ContainerPort,
				HostPort:      &port.HostPort,
			})
		portResources = append(portResources, &mesosproto.Value_Range{
				Begin: proto.Uint64(uint64(port.HostPort)),
				End:   proto.Uint64(uint64(port.HostPort)),
			})
	}
	if len(task.Ports) > 0 {
		dockerInfo.Network= mesosproto.ContainerInfo_DockerInfo_BRIDGE.Enum()
	}

	commandInfo := &mesosproto.CommandInfo{
		Shell: proto.Bool(false),
	}
	if len(task.Arguments) > 0 {
		for _, argument := range task.Arguments {
			commandInfo.Arguments = append(commandInfo.Arguments, argument)
		}
	}
	resources := [] *mesosproto.Resource {
		&mesosproto.Resource{
			Name: proto.String("mem"),
			Scalar: &mesosproto.Value_Scalar{Value: &memory},
			Type: mesosproto.Value_SCALAR.Enum(),
		},
		&mesosproto.Resource{
			Name: proto.String("cpus"),
			Scalar: &mesosproto.Value_Scalar{Value: &cpus},
			Type: mesosproto.Value_SCALAR.Enum(),
		},
	}
	if len(task.Ports) > 0 {
		resources = append(resources,
			&mesosproto.Resource{
				Name: proto.String("ports"),
				Ranges: &mesosproto.Value_Ranges{ Range: portResources},
				Type: mesosproto.Value_RANGES.Enum(),
			},
		)
	}

	taskInfo := &mesosproto.TaskInfo {
		Name: &id,
		TaskId: &mesosproto.TaskID{Value: &id},
		Container: containerInfo,
		Command: commandInfo,
		Resources: resources,
	}

	task.InternalID = id
	task.CreatedAt = time.Now().Unix()
	task.TaskInfo = taskInfo
	task.RequestSent = false

	glog.Infoln("Task: ", task)

	taskRegistry.Register(id, task)

}

func (self *manager) announceFramework() {
	message := &mesosproto.RegisterFrameworkMessage{
		Framework: self.frameworkInfo,
	}

	glog.Infof("Registering with master %s [%s] ", self.masterUPID, message)
	messagePackage := communication.NewMessage(self.masterUPID, message, nil)
	if err := communication.SendMessageToMesos(self.selfUPID, messagePackage); err != nil {
		glog.Errorf("Failed to send RegisterFramework message: %v\n", err)
	}

}

func (self *manager) AcceptOffer(offerId *mesosproto.OfferID, slaveId *mesosproto.SlaveID, taskRequest *managerInterface.Task) {
	glog.Infoln("Working on: ", taskRequest.TaskInfo)
	taskRequest.TaskInfo.SlaveId= slaveId
	message := &mesosproto.LaunchTasksMessage{
		FrameworkId: &mesosproto.FrameworkID{Value: &self.frameworkId},
		OfferIds:    []*mesosproto.OfferID{offerId},
		Tasks:       []*mesosproto.TaskInfo{taskRequest.TaskInfo},
		Filters:     &mesosproto.Filters{},
	}

	messagePackage := communication.NewMessage(self.masterUPID, message, nil)
	if err := communication.SendMessageToMesos(self.selfUPID, messagePackage); err != nil {
		glog.Errorf("Failed to send AcceptOffer message: %v\n", err)
	} else {
		taskRequest.RequestSent = true
	}

}

func (self *manager) DeclineOffer(offerId *mesosproto.OfferID) {
	message := &mesosproto.LaunchTasksMessage{
		FrameworkId: &mesosproto.FrameworkID{Value: &self.frameworkId},
		OfferIds:    []*mesosproto.OfferID{offerId},
		Tasks:       []*mesosproto.TaskInfo{},
		Filters:     &mesosproto.Filters{},
	}

	messagePackage := communication.NewMessage(self.masterUPID, message, nil)
	if err := communication.SendMessageToMesos(self.selfUPID, messagePackage); err != nil {
		glog.Errorf("Failed to send DeclineOffer message: %v\n", err)
	}

}

func (self *manager) acknowledgeStatusUpdate(statusUpdate *mesosproto.StatusUpdateMessage) {
	message := &mesosproto.StatusUpdateAcknowledgementMessage{
		FrameworkId: statusUpdate.GetUpdate().FrameworkId,
		SlaveId:     statusUpdate.GetUpdate().Status.SlaveId,
		TaskId:      statusUpdate.GetUpdate().Status.TaskId,
		Uuid:        statusUpdate.GetUpdate().Uuid,
	}

	messagePackage := communication.NewMessage(self.masterUPID, message, nil)
	if err := communication.SendMessageToMesos(self.selfUPID, messagePackage); err != nil {
		glog.Errorf("Failed to send StatusAccept message: %v\n", err)
	}

}
<commit_msg>Added initial retry logic for mesos connection<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2014 AT&T
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package impl

import (
	"errors"
	"os"
	"fmt"
	"time"
	"strings"

	"code.google.com/p/gogoprotobuf/proto"
	"github.com/golang/glog"

	"github.com/att-innovate/charmander-scheduler/communication"
	"github.com/att-innovate/charmander-scheduler/mesosproto"
	"github.com/att-innovate/charmander-scheduler/scheduler"
	"github.com/att-innovate/charmander-scheduler/upid"

	managerInterface "github.com/att-innovate/charmander-scheduler/manager"
)

var taskRegistry = NewTaskRegistry()
var nodeRegistry = NewNodeRegistry()

func New(
	scheduler *scheduler.Scheduler,
	frameworkName string,
	master string,
	listening string) (manager, error) {

	framework := &mesosproto.FrameworkInfo{
		User: proto.String(""),
		Name: proto.String(frameworkName),
	}

	newManager := manager{
		scheduler:     scheduler,
		frameworkInfo: framework,
		master:        master,
		listening:     listening,
	}

	return newManager, nil
}

type manager struct {
	scheduler     *scheduler.Scheduler
	frameworkInfo *mesosproto.FrameworkInfo
	masterUPID    *upid.UPID
	selfUPID      *upid.UPID
	frameworkId   string
	master        string
	listening     string
}

func (self *manager) Start() error {
	retryCounter := 6

	for ; retryCounter > 0; retryCounter-- {
		if communication.MesosMasterReachable(self.master) { break }
		time.Sleep(10 * time.Second)
	}

	if retryCounter == 0 {
		return errors.New("Mesos unreachable")
	}

	self.frameworkInfo.User = proto.String("root")

	// set default hostname
	if self.frameworkInfo.GetHostname() == "" {
		host, err := os.Hostname()
		if err != nil || host == "" {
			host = "unknown"
		}
		self.frameworkInfo.Hostname = proto.String(host)
	}

	if m, err := upid.Parse("master@" + self.master); err != nil {
		return err
	} else {
		self.masterUPID = m
	}

	self.selfUPID= &upid.UPID{
		ID: "scheduler",
		Host: self.GetListenerIP(),
		Port: fmt.Sprintf("%d", self.GetListenerPortForScheduler())}

	communication.InitRestHandler(self)

	self.announceFramework()

	return nil
}

func (self *manager) GetListenerIP() string {
	return self.listening
}

func (self *manager) GetListenerPortForScheduler() int {
	return 7070
}

func (self *manager) GetListenerPortForRESTApi() int {
	return 7075
}

func (self *manager) GetTaskRequests() []*managerInterface.Task {
	return taskRegistry.Tasks()
}

func (self *manager) ResourceRequirementsWouldMatch(offer *mesosproto.Offer, taskRequest *managerInterface.Task) bool {
	for _, resource := range offer.GetResources() {
		switch {
		case *resource.Name == "cpus":
			if resource.Scalar.GetValue() < float64(taskRequest.Cpus) { return false }
		case *resource.Name == "mem":
			if resource.Scalar.GetValue() < float64(taskRequest.Mem) { return false }
		}
	}
	for _, attribute := range offer.GetAttributes() {
		switch {
		case *attribute.Name == "nodetype":
			if len(taskRequest.NodeType) > 0 && attribute.Text.GetValue() != taskRequest.NodeType { return false }
		case *attribute.Name == "nodename":
			if len(taskRequest.NodeName) > 0 && attribute.Text.GetValue() != taskRequest.NodeName { return false }
		}
	}

	return true
}

func (self *manager) HandleFrameworkRegistered(frameworkId string) {
	self.frameworkId = frameworkId
	if self.scheduler != nil && self.scheduler.Registered != nil {
		self.scheduler.Registered(self, frameworkId)
	}
}

func (self *manager) HandleResourceOffered(offers []*mesosproto.Offer) {
	self.updateNodeRegistry(offers)

	offers= self.enforceSLAs(offers)

	if self.scheduler != nil && self.scheduler.ResourceOffers != nil {
		self.scheduler.ResourceOffers(self, offers)
	}
}

func (self *manager) updateNodeRegistry(offers []*mesosproto.Offer) {
	for _, offer := range offers {
		slaveID := offer.GetSlaveId().GetValue()
		if nodeRegistry.Exists(slaveID) {
			nodeRegistry.UpdateTimeOfLastOffer(slaveID)
			continue
		}

		node := &managerInterface.Node {
			ID: slaveID,
			Hostname: offer.GetHostname(),
			TimeOfLastOffer: time.Now().Unix(),
		}

		for _, attribute := range offer.GetAttributes() {
			if *attribute.Name == "nodename" {
				node.NodeName = attribute.Text.GetValue()
			} else if *attribute.Name == "nodetype" {
				node.NodeType = attribute.Text.GetValue()
			}
		}

		nodeRegistry.Register(slaveID, node)

		glog.Infoln("New Node Registered: ", node)
	}
}

func (self *manager) enforceSLAs(offers []*mesosproto.Offer) []*mesosproto.Offer {
	result := offers
	var taskRequests []*managerInterface.Task
	taskRequests = self.GetTaskRequests()

	for _, taskRequest := range taskRequests {
		if taskRequest.Running { continue }
		if len(taskRequest.Sla) == 0 {continue }

		if taskRequest.Sla == managerInterface.SLA_ONE_PER_NODE {
			for _, offer := range offers {
				if resolveNodeName(*offer) == taskRequest.NodeName {
					if self.ResourceRequirementsWouldMatch(offer, taskRequest) {
						self.AcceptOffer(offer.GetId(), offer.SlaveId, taskRequest)
						result = removeOfferFromList(result, offer)
					}
				}
			}
		}
	}

	return result
}

func resolveNodeName(offer mesosproto.Offer) string {
	for _, attribute := range offer.GetAttributes() {
		if *attribute.Name == "nodename" {
			return attribute.Text.GetValue()
		}
	}

	return "notfound"
}

func removeOfferFromList(offers []*mesosproto.Offer, offer *mesosproto.Offer) []*mesosproto.Offer {
	numberOfOffers := len(offers)
	if numberOfOffers == 1 { return []*mesosproto.Offer {}}

	result := make([]*mesosproto.Offer, numberOfOffers-1)
	i :=  0
	for _, value := range offers {
		if offer.GetId().GetValue() == value.GetId().GetValue() { continue }
		result[i] = value
		i++
	}

	return result
}

func (self *manager) HandleStatusMessage(statusMessage *mesosproto.StatusUpdateMessage) {
	glog.Infof("Status Update %v\n", statusMessage)
	status := statusMessage.GetUpdate().GetStatus()

	switch {
	case  *status.State == mesosproto.TaskState_TASK_RUNNING:
		task, _ := taskRegistry.Fetch(status.GetTaskId().GetValue())
		task.Running= true
		task.SlaveID= status.GetSlaveId().GetValue()
		glog.Infof("Running: %s on Slave: %s", task.InternalID, task.SlaveID)
	case  *status.State == mesosproto.TaskState_TASK_FAILED:
		taskRegistry.Delete(status.GetTaskId().GetValue())
		glog.Infoln("Task Failed: ", status.GetTaskId().GetValue())
	case  *status.State == mesosproto.TaskState_TASK_LOST:
		taskRegistry.Delete(status.GetTaskId().GetValue())
		glog.Infoln("Task Lost: ", status.GetTaskId().GetValue())
	case  *status.State == mesosproto.TaskState_TASK_FINISHED:
		taskRegistry.Delete(status.GetTaskId().GetValue())
		glog.Infoln("Task Finished: ", status.GetTaskId().GetValue())
	}

	self.acknowledgeStatusUpdate(statusMessage)
}

func (self *manager) HandleRunDockerImage(task *managerInterface.Task) {
	if task.Sla == managerInterface.SLA_ONE_PER_NODE {
		for _, node := range nodeRegistry.Nodes() {
			newTask := managerInterface.CopyTask(*task)
			newTask.NodeName= node.NodeName
			self.handleRunDockerImageImpl(&newTask)
		}
	} else {
		self.handleRunDockerImageImpl(task)
	}
}

func (self *manager) handleRunDockerImageImpl(task *managerInterface.Task) {
	id := fmt.Sprintf("%v-%v", strings.Replace(task.ID, " ", "", -1), time.Now().UnixNano())
	memory := float64(task.Mem)
	portResources := []*mesosproto.Value_Range{}
	cpus := float64(0.1)
	if task.Cpus > 0 {
		cpus= task.Cpus
	}

	dockerInfo := &mesosproto.ContainerInfo_DockerInfo {
		Image: &task.DockerImage,
	}
	containerInfo := &mesosproto.ContainerInfo {
		Type: mesosproto.ContainerInfo_DOCKER.Enum(),
		Docker: dockerInfo,
	}
	for _, volume := range task.Volumes {
		mode := mesosproto.Volume_RW
		if volume.Mode == "ro" {
			mode = mesosproto.Volume_RO
		}

		containerInfo.Volumes = append(containerInfo.Volumes, &mesosproto.Volume{
				ContainerPath: &volume.ContainerPath,
				HostPath:      &volume.HostPath,
				Mode:          &mode,
			})
	}
	for _, port := range task.Ports {
		dockerInfo.PortMappings = append(dockerInfo.PortMappings, &mesosproto.ContainerInfo_DockerInfo_PortMapping{
				ContainerPort: &port.ContainerPort,
				HostPort:      &port.HostPort,
			})
		portResources = append(portResources, &mesosproto.Value_Range{
				Begin: proto.Uint64(uint64(port.HostPort)),
				End:   proto.Uint64(uint64(port.HostPort)),
			})
	}
	if len(task.Ports) > 0 {
		dockerInfo.Network= mesosproto.ContainerInfo_DockerInfo_BRIDGE.Enum()
	}

	commandInfo := &mesosproto.CommandInfo{
		Shell: proto.Bool(false),
	}
	if len(task.Arguments) > 0 {
		for _, argument := range task.Arguments {
			commandInfo.Arguments = append(commandInfo.Arguments, argument)
		}
	}
	resources := [] *mesosproto.Resource {
		&mesosproto.Resource{
			Name: proto.String("mem"),
			Scalar: &mesosproto.Value_Scalar{Value: &memory},
			Type: mesosproto.Value_SCALAR.Enum(),
		},
		&mesosproto.Resource{
			Name: proto.String("cpus"),
			Scalar: &mesosproto.Value_Scalar{Value: &cpus},
			Type: mesosproto.Value_SCALAR.Enum(),
		},
	}
	if len(task.Ports) > 0 {
		resources = append(resources,
			&mesosproto.Resource{
				Name: proto.String("ports"),
				Ranges: &mesosproto.Value_Ranges{ Range: portResources},
				Type: mesosproto.Value_RANGES.Enum(),
			},
		)
	}

	taskInfo := &mesosproto.TaskInfo {
		Name: &id,
		TaskId: &mesosproto.TaskID{Value: &id},
		Container: containerInfo,
		Command: commandInfo,
		Resources: resources,
	}

	task.InternalID = id
	task.CreatedAt = time.Now().Unix()
	task.TaskInfo = taskInfo
	task.RequestSent = false

	glog.Infoln("Task: ", task)

	taskRegistry.Register(id, task)

}

func (self *manager) announceFramework() {
	message := &mesosproto.RegisterFrameworkMessage{
		Framework: self.frameworkInfo,
	}

	glog.Infof("Registering with master %s [%s] ", self.masterUPID, message)
	messagePackage := communication.NewMessage(self.masterUPID, message, nil)
	if err := communication.SendMessageToMesos(self.selfUPID, messagePackage); err != nil {
		glog.Errorf("Failed to send RegisterFramework message: %v\n", err)
	}

}

func (self *manager) AcceptOffer(offerId *mesosproto.OfferID, slaveId *mesosproto.SlaveID, taskRequest *managerInterface.Task) {
	glog.Infoln("Working on: ", taskRequest.TaskInfo)
	taskRequest.TaskInfo.SlaveId= slaveId
	message := &mesosproto.LaunchTasksMessage{
		FrameworkId: &mesosproto.FrameworkID{Value: &self.frameworkId},
		OfferIds:    []*mesosproto.OfferID{offerId},
		Tasks:       []*mesosproto.TaskInfo{taskRequest.TaskInfo},
		Filters:     &mesosproto.Filters{},
	}

	messagePackage := communication.NewMessage(self.masterUPID, message, nil)
	if err := communication.SendMessageToMesos(self.selfUPID, messagePackage); err != nil {
		glog.Errorf("Failed to send AcceptOffer message: %v\n", err)
	} else {
		taskRequest.RequestSent = true
	}

}

func (self *manager) DeclineOffer(offerId *mesosproto.OfferID) {
	message := &mesosproto.LaunchTasksMessage{
		FrameworkId: &mesosproto.FrameworkID{Value: &self.frameworkId},
		OfferIds:    []*mesosproto.OfferID{offerId},
		Tasks:       []*mesosproto.TaskInfo{},
		Filters:     &mesosproto.Filters{},
	}

	messagePackage := communication.NewMessage(self.masterUPID, message, nil)
	if err := communication.SendMessageToMesos(self.selfUPID, messagePackage); err != nil {
		glog.Errorf("Failed to send DeclineOffer message: %v\n", err)
	}

}

func (self *manager) acknowledgeStatusUpdate(statusUpdate *mesosproto.StatusUpdateMessage) {
	message := &mesosproto.StatusUpdateAcknowledgementMessage{
		FrameworkId: statusUpdate.GetUpdate().FrameworkId,
		SlaveId:     statusUpdate.GetUpdate().Status.SlaveId,
		TaskId:      statusUpdate.GetUpdate().Status.TaskId,
		Uuid:        statusUpdate.GetUpdate().Uuid,
	}

	messagePackage := communication.NewMessage(self.masterUPID, message, nil)
	if err := communication.SendMessageToMesos(self.selfUPID, messagePackage); err != nil {
		glog.Errorf("Failed to send StatusAccept message: %v\n", err)
	}

}
<|endoftext|>


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

package blueflood

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/square/metrics/api"
	"github.com/square/metrics/log"
)

type httpClient interface {
	// our own client to mock out the standard golang HTTP Client.
	Get(url string) (resp *http.Response, err error)
}

type Config struct {
	BaseUrl  string           `yaml:"base_url"`
	TenantId string           `yaml:"tenant_id"`
	Ttls     map[string]int64 `yaml:"ttls"` // Ttl in days
	Timeout  time.Duration    `yaml:"timeout"`
}

func (c Config) getTTL(r Resolution) time.Duration {
	var ttl int64
	if v, ok := c.Ttls[r.bluefloodEnum]; ok {
		ttl = v
	} else {
		// Use blueflood defaults
		switch r {
		case ResolutionFull:
			ttl = 1
		case Resolution5Min:
			ttl = 30
		case Resolution20Min:
			ttl = 60
		case Resolution60Min:
			ttl = 90
		case Resolution240Min:
			ttl = 180
		case Resolution1440Min:
			ttl = 365
		default:
			// Not a supported resolution by blueflood. No real way to recover if
			// someone's trying to fetch ttl for an invalid resolution.
			panic(fmt.Sprintf("invalid resolution `%s`", r))
		}
	}

	return time.Duration(ttl) * 24 * time.Hour
}

type blueflood struct {
	config Config
	client httpClient
}

type queryResponse struct {
	Values []metricPoint `json:"values"`
}

type metricPoint struct {
	Points    int     `json:"numPoints"`
	Timestamp int64   `json:"timestamp"`
	Average   float64 `json:"average"`
	Max       float64 `json:"max"`
	Min       float64 `json:"min"`
	Variance  float64 `json:"variance"`
}

type Resolution struct {
	bluefloodEnum string
	duration      time.Duration
}

var (
	ResolutionFull    Resolution = Resolution{"FULL", time.Second * 30}
	Resolution5Min               = Resolution{"MIN5", time.Minute * 5}
	Resolution20Min              = Resolution{"MIN20", time.Minute * 20}
	Resolution60Min              = Resolution{"MIN60", time.Minute * 60}
	Resolution240Min             = Resolution{"MIN240", time.Minute * 240}
	Resolution1440Min            = Resolution{"MIN1440", time.Minute * 1440}
)
var Resolutions []Resolution = []Resolution{
	ResolutionFull,
	Resolution5Min,
	Resolution20Min,
	Resolution60Min,
	Resolution240Min,
	Resolution1440Min,
}

func NewBlueflood(c Config) api.Backend {
	b := blueflood{config: c, client: http.DefaultClient}
	b.config.Ttls = map[string]int64{}
	for k, v := range c.Ttls {
		b.config.Ttls[k] = v
	}
	return &b
}

type sampler struct {
	fieldName     string
	fieldSelector func(point metricPoint) float64
	bucketSampler func([]float64) float64
}

// The amount of time before other resolutions become available
// It's not "const" so that it can be mocked out for tests
var availableOnlyFull = time.Hour * 4

func (b *blueflood) FetchSingleSeries(request api.FetchSeriesRequest) (api.Timeseries, error) {
	sampler, ok := samplerMap[request.SampleMethod]
	if !ok {
		return api.Timeseries{}, fmt.Errorf("unsupported SampleMethod %s", request.SampleMethod.String())
	}
	queryResolution := b.config.bluefloodResolution(
		request.Timerange.Resolution(),
		request.Timerange.Start(),
	)

	// Sample the data at the given `queryResolution`
	queryUrl, err := b.constructURL(request, sampler, queryResolution)
	if err != nil {
		return api.Timeseries{}, err
	}
	parsedResult, err := b.fetch(request, queryUrl)
	if err != nil {
		return api.Timeseries{}, err
	}

	combinedResult := parsedResult.Values

	// Sample the data at the FULL resolution.
	// In order to do this, we use the same `request` object but hard-code the ResolutionFull parameter.
	fullResolutionQueryURL, err := b.constructURL(request, sampler, ResolutionFull)
	if err == nil {
		fullResolutionParsedResult, err := b.fetch(request, fullResolutionQueryURL)
		if err == nil {
			combinedResult = append(parsedResult.Values, fullResolutionParsedResult.Values...)
		}
	}

	values := processResult(combinedResult, request.Timerange, sampler, queryResolution)
	log.Debugf("Constructed timeseries from result: %v", values)

	return api.Timeseries{
		Values: values,
		TagSet: request.Metric.TagSet,
	}, nil
}

// Helper functions
// ----------------

// constructURL creates the URL to the blueflood's backend to fetch the data from.
func (b *blueflood) constructURL(
	request api.FetchSeriesRequest,
	sampler sampler,
	queryResolution Resolution,
) (*url.URL, error) {
	graphiteName, err := request.API.ToGraphiteName(request.Metric)
	if err != nil {
		return nil, api.BackendError{request.Metric, api.InvalidSeriesError, "cannot convert to graphite name"}
	}

	result, err := url.Parse(fmt.Sprintf("%s/v2.0/%s/views/%s", b.config.BaseUrl, b.config.TenantId, graphiteName))
	if err != nil {
		return nil, api.BackendError{request.Metric, api.InvalidSeriesError, "cannot generate URL"}
	}

	params := url.Values{}
	params.Set("from", strconv.FormatInt(request.Timerange.Start(), 10))
	// Pull a bit outside of the requested range from blueflood so we
	// have enough data to generate all snapped values
	params.Set("to", strconv.FormatInt(request.Timerange.End()+request.Timerange.ResolutionMillis(), 10))
	params.Set("resolution", queryResolution.bluefloodEnum)
	params.Set("select", fmt.Sprintf("numPoints,%s", strings.ToLower(sampler.fieldName)))
	result.RawQuery = params.Encode()
	return result, nil
}

// fetches from the backend. on error, it returns an instance of api.BackendError
func (b *blueflood) fetch(request api.FetchSeriesRequest, queryUrl *url.URL) (queryResponse, error) {
	log.Debugf("Blueflood fetch: %s", queryUrl.String())
	success := make(chan queryResponse)
	failure := make(chan error)
	timeout := time.After(b.config.Timeout)
	go func() {
		resp, err := b.client.Get(queryUrl.String())
		if err != nil {
			failure <- api.BackendError{request.Metric, api.FetchIOError, "error while fetching - http connection"}
			return
		}
		defer resp.Body.Close()

		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			failure <- api.BackendError{request.Metric, api.FetchIOError, "error while fetching - reading"}
			return
		}

		log.Debugf("Fetch result: %s", string(body))

		var parsedJson queryResponse
		err = json.Unmarshal(body, &parsedJson)
		// Construct a Timeseries from the result:
		if err != nil {
			failure <- api.BackendError{request.Metric, api.FetchIOError, "error while fetching - json decoding"}
			return
		}
		success <- parsedJson
	}()
	select {
	case response := <-success:
		return response, nil
	case err := <-failure:
		return queryResponse{}, err
	case <-timeout:
		return queryResponse{}, api.BackendError{request.Metric, api.FetchTimeoutError, ""}
	}
}

func processResult(
	points []metricPoint,
	timerange api.Timerange,
	sampler sampler,
	queryResolution Resolution) []float64 {
	// buckets are each filled with from the points stored in `points`, according to their timestamps.
	buckets := bucketsFromMetricPoints(points, sampler.fieldSelector, timerange)

	// values will hold the final values to be returned as the series.
	values := make([]float64, timerange.Slots())

	for i, bucket := range buckets {
		if len(bucket) == 0 {
			values[i] = math.NaN()
			continue
		}
		values[i] = sampler.bucketSampler(bucket)
	}

	// interpolate.
	return values
}

func addMetricPoint(metricPoint metricPoint, field func(metricPoint) float64, timerange api.Timerange, buckets [][]float64) bool {
	value := field(metricPoint)
	// The index to assign within the array is computed using the timestamp.
	// It floors to the nearest index.
	index := (metricPoint.Timestamp - timerange.Start()) / timerange.ResolutionMillis()
	if index < 0 || index >= int64(timerange.Slots()) {
		return false
	}
	buckets[index] = append(buckets[index], value)
	return true
}

func bucketsFromMetricPoints(metricPoints []metricPoint, resultField func(metricPoint) float64, timerange api.Timerange) [][]float64 {
	buckets := make([][]float64, timerange.Slots())
	for _, point := range metricPoints {
		addMetricPoint(point, resultField, timerange, buckets)
	}
	return buckets
}

var samplerMap map[api.SampleMethod]sampler = map[api.SampleMethod]sampler{
	api.SampleMean: {
		fieldName:     "average",
		fieldSelector: func(point metricPoint) float64 { return point.Average },
		bucketSampler: func(bucket []float64) float64 {
			value := 0.0
			count := 0.0
			for _, v := range bucket {
				if !math.IsNaN(v) {
					value += v
					count++
				}
			}
			return value / count
		},
	},
	api.SampleMin: {
		fieldName:     "min",
		fieldSelector: func(point metricPoint) float64 { return point.Min },
		bucketSampler: func(bucket []float64) float64 {
			value := bucket[0]
			for _, v := range bucket {
				if math.IsNaN(v) {
					continue
				}
				if math.IsNaN(value) {
					value = v
				}
				value = math.Min(value, v)
			}
			return value
		},
	},
	api.SampleMax: {
		fieldName:     "max",
		fieldSelector: func(point metricPoint) float64 { return point.Max },
		bucketSampler: func(bucket []float64) float64 {
			value := bucket[0]
			for _, v := range bucket {
				if math.IsNaN(v) {
					continue
				}
				if math.IsNaN(value) {
					value = v
				}
				value = math.Max(value, v)
			}
			return value
		},
	},
}

// Blueflood keys the resolution param to a java enum, so we have to convert
// between them.
func (c Config) bluefloodResolution(
	desiredResolution time.Duration,
	startMs int64) Resolution {
	now := time.Now().Unix() * 1000
	// Choose the appropriate resolution based on TTL, fetching the highest resolution data we can
	for _, current := range Resolutions {
		age := time.Duration(now-startMs) * time.Millisecond
		maxAge := c.getTTL(current)
		log.Debugf("Desired (s): %d\n", desiredResolution/time.Second)
		log.Debugf("Current (s): %d\n", current.duration/time.Second)
		log.Debugf("age (s): %d\n", age/time.Second)
		log.Debugf("ttl (s): %d\n", maxAge/time.Second)
		if desiredResolution <= current.duration &&
			age < c.getTTL(current) {
			return current
		}
	}
	// return the coarsest resolution.
	return Resolutions[len(Resolutions)-1]
}
<commit_msg>fix up the sampling methods<commit_after>// Copyright 2015 Square Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package blueflood

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/square/metrics/api"
	"github.com/square/metrics/log"
)

type httpClient interface {
	// our own client to mock out the standard golang HTTP Client.
	Get(url string) (resp *http.Response, err error)
}

type Config struct {
	BaseUrl  string           `yaml:"base_url"`
	TenantId string           `yaml:"tenant_id"`
	Ttls     map[string]int64 `yaml:"ttls"` // Ttl in days
	Timeout  time.Duration    `yaml:"timeout"`
}

func (c Config) getTTL(r Resolution) time.Duration {
	var ttl int64
	if v, ok := c.Ttls[r.bluefloodEnum]; ok {
		ttl = v
	} else {
		// Use blueflood defaults
		switch r {
		case ResolutionFull:
			ttl = 1
		case Resolution5Min:
			ttl = 30
		case Resolution20Min:
			ttl = 60
		case Resolution60Min:
			ttl = 90
		case Resolution240Min:
			ttl = 180
		case Resolution1440Min:
			ttl = 365
		default:
			// Not a supported resolution by blueflood. No real way to recover if
			// someone's trying to fetch ttl for an invalid resolution.
			panic(fmt.Sprintf("invalid resolution `%s`", r))
		}
	}

	return time.Duration(ttl) * 24 * time.Hour
}

type blueflood struct {
	config Config
	client httpClient
}

type queryResponse struct {
	Values []metricPoint `json:"values"`
}

type metricPoint struct {
	Points    int     `json:"numPoints"`
	Timestamp int64   `json:"timestamp"`
	Average   float64 `json:"average"`
	Max       float64 `json:"max"`
	Min       float64 `json:"min"`
	Variance  float64 `json:"variance"`
}

type Resolution struct {
	bluefloodEnum string
	duration      time.Duration
}

var (
	ResolutionFull    Resolution = Resolution{"FULL", time.Second * 30}
	Resolution5Min               = Resolution{"MIN5", time.Minute * 5}
	Resolution20Min              = Resolution{"MIN20", time.Minute * 20}
	Resolution60Min              = Resolution{"MIN60", time.Minute * 60}
	Resolution240Min             = Resolution{"MIN240", time.Minute * 240}
	Resolution1440Min            = Resolution{"MIN1440", time.Minute * 1440}
)
var Resolutions []Resolution = []Resolution{
	ResolutionFull,
	Resolution5Min,
	Resolution20Min,
	Resolution60Min,
	Resolution240Min,
	Resolution1440Min,
}

func NewBlueflood(c Config) api.Backend {
	b := blueflood{config: c, client: http.DefaultClient}
	b.config.Ttls = map[string]int64{}
	for k, v := range c.Ttls {
		b.config.Ttls[k] = v
	}
	return &b
}

type sampler struct {
	fieldName     string
	fieldSelector func(point metricPoint) float64
	bucketSampler func([]float64) float64
}

// The amount of time before other resolutions become available
// It's not "const" so that it can be mocked out for tests
var availableOnlyFull = time.Hour * 4

func (b *blueflood) FetchSingleSeries(request api.FetchSeriesRequest) (api.Timeseries, error) {
	sampler, ok := samplerMap[request.SampleMethod]
	if !ok {
		return api.Timeseries{}, fmt.Errorf("unsupported SampleMethod %s", request.SampleMethod.String())
	}
	queryResolution := b.config.bluefloodResolution(
		request.Timerange.Resolution(),
		request.Timerange.Start(),
	)

	// Sample the data at the given `queryResolution`
	queryUrl, err := b.constructURL(request, sampler, queryResolution)
	if err != nil {
		return api.Timeseries{}, err
	}
	parsedResult, err := b.fetch(request, queryUrl)
	if err != nil {
		return api.Timeseries{}, err
	}

	combinedResult := parsedResult.Values

	// Sample the data at the FULL resolution.
	// In order to do this, we use the same `request` object but hard-code the ResolutionFull parameter.
	fullResolutionQueryURL, err := b.constructURL(request, sampler, ResolutionFull)
	if err == nil {
		fullResolutionParsedResult, err := b.fetch(request, fullResolutionQueryURL)
		if err == nil {
			combinedResult = append(parsedResult.Values, fullResolutionParsedResult.Values...)
		}
	}

	values := processResult(combinedResult, request.Timerange, sampler, queryResolution)
	log.Debugf("Constructed timeseries from result: %v", values)

	return api.Timeseries{
		Values: values,
		TagSet: request.Metric.TagSet,
	}, nil
}

// Helper functions
// ----------------

// constructURL creates the URL to the blueflood's backend to fetch the data from.
func (b *blueflood) constructURL(
	request api.FetchSeriesRequest,
	sampler sampler,
	queryResolution Resolution,
) (*url.URL, error) {
	graphiteName, err := request.API.ToGraphiteName(request.Metric)
	if err != nil {
		return nil, api.BackendError{request.Metric, api.InvalidSeriesError, "cannot convert to graphite name"}
	}

	result, err := url.Parse(fmt.Sprintf("%s/v2.0/%s/views/%s", b.config.BaseUrl, b.config.TenantId, graphiteName))
	if err != nil {
		return nil, api.BackendError{request.Metric, api.InvalidSeriesError, "cannot generate URL"}
	}

	params := url.Values{}
	params.Set("from", strconv.FormatInt(request.Timerange.Start(), 10))
	// Pull a bit outside of the requested range from blueflood so we
	// have enough data to generate all snapped values
	params.Set("to", strconv.FormatInt(request.Timerange.End()+request.Timerange.ResolutionMillis(), 10))
	params.Set("resolution", queryResolution.bluefloodEnum)
	params.Set("select", fmt.Sprintf("numPoints,%s", strings.ToLower(sampler.fieldName)))
	result.RawQuery = params.Encode()
	return result, nil
}

// fetches from the backend. on error, it returns an instance of api.BackendError
func (b *blueflood) fetch(request api.FetchSeriesRequest, queryUrl *url.URL) (queryResponse, error) {
	log.Debugf("Blueflood fetch: %s", queryUrl.String())
	success := make(chan queryResponse)
	failure := make(chan error)
	timeout := time.After(b.config.Timeout)
	go func() {
		resp, err := b.client.Get(queryUrl.String())
		if err != nil {
			failure <- api.BackendError{request.Metric, api.FetchIOError, "error while fetching - http connection"}
			return
		}
		defer resp.Body.Close()

		body, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			failure <- api.BackendError{request.Metric, api.FetchIOError, "error while fetching - reading"}
			return
		}

		log.Debugf("Fetch result: %s", string(body))

		var parsedJson queryResponse
		err = json.Unmarshal(body, &parsedJson)
		// Construct a Timeseries from the result:
		if err != nil {
			failure <- api.BackendError{request.Metric, api.FetchIOError, "error while fetching - json decoding"}
			return
		}
		success <- parsedJson
	}()
	select {
	case response := <-success:
		return response, nil
	case err := <-failure:
		return queryResponse{}, err
	case <-timeout:
		return queryResponse{}, api.BackendError{request.Metric, api.FetchTimeoutError, ""}
	}
}

func processResult(
	points []metricPoint,
	timerange api.Timerange,
	sampler sampler,
	queryResolution Resolution) []float64 {
	// buckets are each filled with from the points stored in `points`, according to their timestamps.
	buckets := bucketsFromMetricPoints(points, sampler.fieldSelector, timerange)

	// values will hold the final values to be returned as the series.
	values := make([]float64, timerange.Slots())

	for i, bucket := range buckets {
		if len(bucket) == 0 {
			values[i] = math.NaN()
			continue
		}
		values[i] = sampler.bucketSampler(bucket)
	}

	// interpolate.
	return values
}

func addMetricPoint(metricPoint metricPoint, field func(metricPoint) float64, timerange api.Timerange, buckets [][]float64) bool {
	value := field(metricPoint)
	// The index to assign within the array is computed using the timestamp.
	// It floors to the nearest index.
	index := (metricPoint.Timestamp - timerange.Start()) / timerange.ResolutionMillis()
	if index < 0 || index >= int64(timerange.Slots()) {
		return false
	}
	buckets[index] = append(buckets[index], value)
	return true
}

func bucketsFromMetricPoints(metricPoints []metricPoint, resultField func(metricPoint) float64, timerange api.Timerange) [][]float64 {
	buckets := make([][]float64, timerange.Slots())
	for _, point := range metricPoints {
		addMetricPoint(point, resultField, timerange, buckets)
	}
	return buckets
}

var samplerMap map[api.SampleMethod]sampler = map[api.SampleMethod]sampler{
	api.SampleMean: {
		fieldName:     "average",
		fieldSelector: func(point metricPoint) float64 { return point.Average },
		bucketSampler: func(bucket []float64) float64 {
			value := 0.0
			count := 0
			for _, v := range bucket {
				if !math.IsNaN(v) {
					value += v
					count++
				}
			}
			return value / float64(count)
		},
	},
	api.SampleMin: {
		fieldName:     "min",
		fieldSelector: func(point metricPoint) float64 { return point.Min },
		bucketSampler: func(bucket []float64) float64 {
			smallest := math.NaN()
			for _, v := range bucket {
				if math.IsNaN(v) {
					continue
				}
				if math.IsNaN(smallest) {
					smallest = v
				} else {
					smallest = math.Min(smallest, v)
				}
			}
			return smallest
		},
	},
	api.SampleMax: {
		fieldName:     "max",
		fieldSelector: func(point metricPoint) float64 { return point.Max },
		bucketSampler: func(bucket []float64) float64 {
			largest := math.NaN()
			for _, v := range bucket {
				if math.IsNaN(v) {
					continue
				}
				if math.IsNaN(largest) {
					largest = v
				} else {
					largest = math.Max(largest, v)
				}
			}
			return largest
		},
	},
}

// Blueflood keys the resolution param to a java enum, so we have to convert
// between them.
func (c Config) bluefloodResolution(
	desiredResolution time.Duration,
	startMs int64) Resolution {
	now := time.Now().Unix() * 1000
	// Choose the appropriate resolution based on TTL, fetching the highest resolution data we can
	for _, current := range Resolutions {
		age := time.Duration(now-startMs) * time.Millisecond
		maxAge := c.getTTL(current)
		log.Debugf("Desired (s): %d\n", desiredResolution/time.Second)
		log.Debugf("Current (s): %d\n", current.duration/time.Second)
		log.Debugf("age (s): %d\n", age/time.Second)
		log.Debugf("ttl (s): %d\n", maxAge/time.Second)
		if desiredResolution <= current.duration &&
			age < c.getTTL(current) {
			return current
		}
	}
	// return the coarsest resolution.
	return Resolutions[len(Resolutions)-1]
}
<|endoftext|>