<commit_before>package goquery

import (
	"errors"
	"io"
	"net/http"
	"net/url"

	"golang.org/x/net/html"
)

// Document represents an HTML document to be manipulated. Unlike jQuery, which
// is loaded as part of a DOM document, and thus acts upon its containing
// document, GoQuery doesn't know which HTML document to act upon. So it needs
// to be told, and that's what the Document class is for. It holds the root
// document node to manipulate, and can make selections on this document.
type Document struct {
	*Selection
	Url      *url.URL
	rootNode *html.Node
}

// NewDocumentFromNode is a Document constructor that takes a root html Node
// as argument.
func NewDocumentFromNode(root *html.Node) *Document {
	return newDocument(root, nil)
}

// NewDocument is a Document constructor that takes a string URL as argument.
// It loads the specified document, parses it, and stores the root Document
// node, ready to be manipulated.
func NewDocument(url string) (*Document, error) {
	// Load the URL
	res, e := http.Get(url)
	if e != nil {
		return nil, e
	}
	return NewDocumentFromResponse(res)
}

// NewDocumentFromReader returns a Document from a generic reader.
// It returns an error as second value if the reader's data cannot be parsed
// as html. It does *not* check if the reader is also an io.Closer, so the
// provided reader is never closed by this call, it is the responsibility
// of the caller to close it if required.
func NewDocumentFromReader(r io.Reader) (*Document, error) {
	root, e := html.Parse(r)
	if e != nil {
		return nil, e
	}
	return newDocument(root, nil), nil
}

// NewDocumentFromResponse is another Document constructor that takes an http resonse as argument.
// It loads the specified response's document, parses it, and stores the root Document
// node, ready to be manipulated. The response's body is closed on return.
func NewDocumentFromResponse(res *http.Response) (*Document, error) {
	if res == nil {
		return nil, errors.New("Response is nil pointer")
	}

	defer res.Body.Close()

	// Parse the HTML into nodes
	root, e := html.Parse(res.Body)
	if e != nil {
		return nil, e
	}

	// Create and fill the document
	return newDocument(root, res.Request.URL), nil
}

// Create a deep-clone of a document
func NewDocumentFromDocument(doc *Document) *Document {
	return newDocument(cloneNode(doc.rootNode), doc.Url)
}

// Private constructor, make sure all fields are correctly filled.
func newDocument(root *html.Node, url *url.URL) *Document {
	// Create and fill the document
	d := &Document{nil, url, root}
	d.Selection = newSingleSelection(root, d)
	return d
}

// Selection represents a collection of nodes matching some criteria. The
// initial Selection can be created by using Document.Find, and then
// manipulated using the jQuery-like chainable syntax and methods.
type Selection struct {
	Nodes    []*html.Node
	document *Document
	prevSel  *Selection
}

// Helper constructor to create an empty selection
func newEmptySelection(doc *Document) *Selection {
	return &Selection{nil, doc, nil}
}

// Helper constructor to create a selection of only one node
func newSingleSelection(node *html.Node, doc *Document) *Selection {
	return &Selection{[]*html.Node{node}, doc, nil}
}

// Matcher is an interface that defines the methods to match
// HTML nodes against a compiled selector string. Cascadia's
// Selector implements this interface.
type Matcher interface {
	Match(*html.Node) bool
	MatchAll(*html.Node) []*html.Node
	Filter([]*html.Node) []*html.Node
}
<commit_msg>rename NewDocumentFromDocument to CloneDocument<commit_after>package goquery

import (
	"errors"
	"io"
	"net/http"
	"net/url"

	"golang.org/x/net/html"
)

// Document represents an HTML document to be manipulated. Unlike jQuery, which
// is loaded as part of a DOM document, and thus acts upon its containing
// document, GoQuery doesn't know which HTML document to act upon. So it needs
// to be told, and that's what the Document class is for. It holds the root
// document node to manipulate, and can make selections on this document.
type Document struct {
	*Selection
	Url      *url.URL
	rootNode *html.Node
}

// NewDocumentFromNode is a Document constructor that takes a root html Node
// as argument.
func NewDocumentFromNode(root *html.Node) *Document {
	return newDocument(root, nil)
}

// NewDocument is a Document constructor that takes a string URL as argument.
// It loads the specified document, parses it, and stores the root Document
// node, ready to be manipulated.
func NewDocument(url string) (*Document, error) {
	// Load the URL
	res, e := http.Get(url)
	if e != nil {
		return nil, e
	}
	return NewDocumentFromResponse(res)
}

// NewDocumentFromReader returns a Document from a generic reader.
// It returns an error as second value if the reader's data cannot be parsed
// as html. It does *not* check if the reader is also an io.Closer, so the
// provided reader is never closed by this call, it is the responsibility
// of the caller to close it if required.
func NewDocumentFromReader(r io.Reader) (*Document, error) {
	root, e := html.Parse(r)
	if e != nil {
		return nil, e
	}
	return newDocument(root, nil), nil
}

// NewDocumentFromResponse is another Document constructor that takes an http resonse as argument.
// It loads the specified response's document, parses it, and stores the root Document
// node, ready to be manipulated. The response's body is closed on return.
func NewDocumentFromResponse(res *http.Response) (*Document, error) {
	if res == nil {
		return nil, errors.New("Response is nil pointer")
	}

	defer res.Body.Close()

	// Parse the HTML into nodes
	root, e := html.Parse(res.Body)
	if e != nil {
		return nil, e
	}

	// Create and fill the document
	return newDocument(root, res.Request.URL), nil
}

// CloneDocument creates a deep-clone of a document.
func CloneDocument(doc *Document) *Document {
	return newDocument(cloneNode(doc.rootNode), doc.Url)
}

// Private constructor, make sure all fields are correctly filled.
func newDocument(root *html.Node, url *url.URL) *Document {
	// Create and fill the document
	d := &Document{nil, url, root}
	d.Selection = newSingleSelection(root, d)
	return d
}

// Selection represents a collection of nodes matching some criteria. The
// initial Selection can be created by using Document.Find, and then
// manipulated using the jQuery-like chainable syntax and methods.
type Selection struct {
	Nodes    []*html.Node
	document *Document
	prevSel  *Selection
}

// Helper constructor to create an empty selection
func newEmptySelection(doc *Document) *Selection {
	return &Selection{nil, doc, nil}
}

// Helper constructor to create a selection of only one node
func newSingleSelection(node *html.Node, doc *Document) *Selection {
	return &Selection{[]*html.Node{node}, doc, nil}
}

// Matcher is an interface that defines the methods to match
// HTML nodes against a compiled selector string. Cascadia's
// Selector implements this interface.
type Matcher interface {
	Match(*html.Node) bool
	MatchAll(*html.Node) []*html.Node
	Filter([]*html.Node) []*html.Node
}
<|endoftext|>


<commit_before>package main

// Just enough UPnP to be able to forward ports
//

import (
    "bytes"
    "jackpal/http"
    "log"
    "os"
    "net"
    "strings"
    "strconv"
    "xml"
)

type upnpNAT struct {
    serviceURL string
}

type NAT interface {
    ForwardPort(protocol string, externalPort, internalPort int, description string, timeout int) (err os.Error)
    DeleteForwardingRule(protocol string, externalPort int) (err os.Error)
}

func Discover() (nat NAT, err os.Error) {
    ssdp, err := net.ResolveUDPAddr("239.255.255.250:1900")
    if err != nil {
        return
    }
    conn, err := net.ListenPacket("udp4", ":0")
    if err != nil {
        return
    }
    socket := conn.(*net.UDPConn)

    err = socket.SetReadTimeout(3 * 1000 * 1000 * 1000)
    if err != nil {
        return
    }

    buf := bytes.NewBufferString(
		"M-SEARCH * HTTP/1.1\r\n" +
		"HOST: 239.255.255.250:1900\r\n" +
		"ST:upnp:rootdevice\r\n" +
		"MAN:\"ssdp:discover\"\r\n" +
		"MX:3\r\n\r\n")
    message := buf.Bytes()
    for i := 0; i < 3; i++ {
        _, err = socket.WriteToUDP(message, ssdp)
		if err != nil {
			return
		}
    }
    answerBytes := make([]byte, 1024)
    for i := 0; i < 10; i++ {
        var n int
    	n, _, err = socket.ReadFromUDP(answerBytes)
		if err != nil {
            socket.Close()
			return
		}
		answer := string(answerBytes[0:n])
        if strings.Index(answer, "\r\nST: upnp:rootdevice\r\n") < 0 {
            continue
        }
        locString := "\r\nLocation: "
        locIndex := strings.Index(answer, locString)
        if locIndex < 0 {
            continue
        }
        loc := answer[locIndex + len(locString):]
        endIndex := strings.Index(loc, "\r\n")
        if endIndex < 0 {
            continue
        }
        locURL := loc[0:endIndex]
        var serviceURL string
        serviceURL, err = getServiceURL(locURL)
        if err != nil {
            return
        }
        nat = &upnpNAT{serviceURL: serviceURL}
        break
    }
    socket.Close()
    return
}

type Service struct {
    ServiceType string
    ControlURL string
}

type DeviceList struct {
    Device []Device
}

type ServiceList struct {
    Service []Service
}

type Device struct {
    DeviceType string
    DeviceList DeviceList
    ServiceList ServiceList
}

type Root struct {
    Device Device
}

func getChildDevice(d *Device, deviceType string) *Device {
    dl := d.DeviceList.Device
    for i := 0; i < len(dl); i++ {
        if dl[i].DeviceType == deviceType {
            return &dl[i]
        }
    }
    return nil
}

func getChildService(d *Device, serviceType string) *Service {
    sl := d.ServiceList.Service
    for i := 0; i < len(sl); i++ {
        if sl[i].ServiceType == serviceType {
            return &sl[i]
        }
    }
    return nil
}

func getServiceURL(rootURL string) (url string, err os.Error) {
    r, _, err := http.Get(rootURL)
	if err != nil {
		return
	}
	defer r.Body.Close()
	if r.StatusCode >= 400 {
        err = os.NewError(string(r.StatusCode))
        return
    }
    var root Root
    err = xml.Unmarshal(r.Body, &root)
	if err != nil {
		return
	}
	a := &root.Device
	if a.DeviceType != "urn:schemas-upnp-org:device:InternetGatewayDevice:1" {
	    err = os.NewError("No InternetGatewayDevice")
	    return
	}
	b := getChildDevice(a, "urn:schemas-upnp-org:device:WANDevice:1")
	if b == nil {
	    err = os.NewError("No WANDevice")
	    return
	}
	c := getChildDevice(b, "urn:schemas-upnp-org:device:WANConnectionDevice:1")
	if c == nil {
	    err = os.NewError("No WANConnectionDevice")
	    return
	}
	d := getChildService(c, "urn:schemas-upnp-org:service:WANIPConnection:1")
	if d == nil {
	    err = os.NewError("No WANIPConnection")
	    return
	}
	url = combineURL(rootURL, d.ControlURL)
	return
}

func combineURL(rootURL, subURL string) string {
    protocolEnd := "://"
    protoEndIndex := strings.Index(rootURL, protocolEnd)
    a := rootURL[protoEndIndex + len(protocolEnd):]
    rootIndex := strings.Index(a, "/")
    return rootURL[0:protoEndIndex + len(protocolEnd) + rootIndex] + subURL
}

type stringBuffer struct {
    base string
    current string
}

func NewStringBuffer(s string) *stringBuffer {
    return &stringBuffer{s, s}
}

func (sb *stringBuffer) Read(p []byte) (n int, err os.Error) {
    s := sb.current
    lenStr := len(s)
    if lenStr == 0 {
        return 0, os.EOF
    }
    n = len(p)
    if n > lenStr {
        n = lenStr
        err = os.EOF
    }
    for i := 0; i < n; i++ {
        p[i] = s[i]
    }
    sb.current = s[n:]
    return
}

func (sb *stringBuffer) Seek(offset int64, whence int) (ret int64, err os.Error) {
    var newOffset int64
    switch whence {
    case 0: // from beginning
        newOffset = offset
    case 1: // relative
        newOffset = int64(len(sb.base) - len(sb.current)) + offset
    case 2: // from end
        newOffset = int64(len(sb.base)) - offset
    default:
        err = os.NewError("bad whence")
        return
    }
    if newOffset < 0 || newOffset > int64(len(sb.base)) {
        err = os.NewError("offset out of range")
    } else {
        sb.current = sb.base[newOffset:]
        ret = newOffset
    }
    return
}

func soapRequest(url, function, message string) (r *http.Response, err os.Error) {
    fullMessage := "<?xml version=\"1.0\"?>" +
        "<s:Envelope xmlns:s=\"http://schemas.xmlsoap.org/soap/envelope/\" s:encodingStyle=\"http://schemas.xmlsoap.org/soap/encoding/\">\r\n" +
        "<s:Body>" + message + "</s:Body></s:Envelope>"

	var req http.Request
	req.Method = "POST"
	req.Body = NewStringBuffer(fullMessage)
	req.UserAgent = "Darwin/10.0.0, UPnP/1.0, MiniUPnPc/1.3"
	req.Header = map[string]string{
		"Content-Type": "text/xml",
		// "Transfer-Encoding": "chunked",
		"SOAPAction": "\"urn:schemas-upnp-org:service:WANIPConnection:1#" + function + "\"",
		"Connection": "Close",
		"Cache-Control": "no-cache",
		"Pragma": "no-cache",
	}

	req.URL, err = http.ParseURL(url)
	if err != nil {
		return
	}

	// log.Stderr("soapRequest ", req)

	r, err = http.Send(&req)

	if r.StatusCode >= 4000 {
        err = os.NewError("Error " + strconv.Itoa(r.StatusCode) + " for " + function)
        r.Body.Close()
        r = nil
        return
    }
	return
}

func (n *upnpNAT) ForwardPort(protocol string, externalPort, internalPort int,
    description string, timeout int) (err os.Error) {
    message := "<u:AddPortMapping xmlns:u=\"urn:schemas-upnp-org:service:WANIPConnection:1\">\r\n" +
        "<NewRemoteHost></NewRemoteHost><NewExternalPort>" + strconv.Itoa(externalPort) +
        "</NewExternalPort><NewProtocol>" + protocol + "</NewProtocol>" +
        "<NewInternalPort>" + strconv.Itoa(internalPort) + "</NewInternalPort><NewInternalClient>" +
        "192.168.0.124" + // TODO: Put our IP address here.
        "</NewInternalClient><NewEnabled>1</NewEnabled><NewPortMappingDescription>" +
        description +
        "</NewPortMappingDescription><NewLeaseDuration>" + strconv.Itoa(timeout) + "</NewLeaseDuration></u:AddPortMapping>"

    var response *http.Response
    response, err = soapRequest(n.serviceURL, "AddPortMapping", message)

    // TODO: check response to see if the port was forwarded
    _ = response
    return
}

func (n *upnpNAT) DeleteForwardingRule(protocol string, externalPort int) (err os.Error) {
    return
}

func testUPnP() {
    log.Stderr("Starting UPnP test")
    nat, err := Discover()
    log.Stderr("nat ", nat, "err ", err)
    if err != nil {
        return
    }
    port := 60001
    err = nat.ForwardPort("TCP", port, port, "Taipei-Torrent", 0)
    log.Stderr("err ", err)
    err = nat.DeleteForwardingRule("TCP", port)
    log.Stderr("err ", err)
}
<commit_msg>Try to receive UDP packets multiple times before giving up.<commit_after>package main

// Just enough UPnP to be able to forward ports
//

import (
    "bytes"
    "jackpal/http"
    "log"
    "os"
    "net"
    "strings"
    "strconv"
    "xml"
)

type upnpNAT struct {
    serviceURL string
}

type NAT interface {
    ForwardPort(protocol string, externalPort, internalPort int, description string, timeout int) (err os.Error)
    DeleteForwardingRule(protocol string, externalPort int) (err os.Error)
}

func Discover() (nat NAT, err os.Error) {
    ssdp, err := net.ResolveUDPAddr("239.255.255.250:1900")
    if err != nil {
        return
    }
    conn, err := net.ListenPacket("udp4", ":0")
    if err != nil {
        return
    }
    socket := conn.(*net.UDPConn)

    err = socket.SetReadTimeout(3 * 1000 * 1000 * 1000)
    if err != nil {
        return
    }

    buf := bytes.NewBufferString(
		"M-SEARCH * HTTP/1.1\r\n" +
		"HOST: 239.255.255.250:1900\r\n" +
		"ST:upnp:rootdevice\r\n" +
		"MAN:\"ssdp:discover\"\r\n" +
		"MX:3\r\n\r\n")
    message := buf.Bytes()
    for i := 0; i < 3; i++ {
        _, err = socket.WriteToUDP(message, ssdp)
		if err != nil {
			return
		}
    }
    answerBytes := make([]byte, 1024)
    for i := 0; i < 10; i++ {
        var n int
    	n, _, err = socket.ReadFromUDP(answerBytes)
		if err != nil {
		    continue
            // socket.Close()
			// return
		}
		answer := string(answerBytes[0:n])
        if strings.Index(answer, "\r\nST: upnp:rootdevice\r\n") < 0 {
            continue
        }
        locString := "\r\nLocation: "
        locIndex := strings.Index(answer, locString)
        if locIndex < 0 {
            continue
        }
        loc := answer[locIndex + len(locString):]
        endIndex := strings.Index(loc, "\r\n")
        if endIndex < 0 {
            continue
        }
        locURL := loc[0:endIndex]
        var serviceURL string
        serviceURL, err = getServiceURL(locURL)
        if err != nil {
            return
        }
        nat = &upnpNAT{serviceURL: serviceURL}
        break
    }
    socket.Close()
    return
}

type Service struct {
    ServiceType string
    ControlURL string
}

type DeviceList struct {
    Device []Device
}

type ServiceList struct {
    Service []Service
}

type Device struct {
    DeviceType string
    DeviceList DeviceList
    ServiceList ServiceList
}

type Root struct {
    Device Device
}

func getChildDevice(d *Device, deviceType string) *Device {
    dl := d.DeviceList.Device
    for i := 0; i < len(dl); i++ {
        if dl[i].DeviceType == deviceType {
            return &dl[i]
        }
    }
    return nil
}

func getChildService(d *Device, serviceType string) *Service {
    sl := d.ServiceList.Service
    for i := 0; i < len(sl); i++ {
        if sl[i].ServiceType == serviceType {
            return &sl[i]
        }
    }
    return nil
}

func getServiceURL(rootURL string) (url string, err os.Error) {
    r, _, err := http.Get(rootURL)
	if err != nil {
		return
	}
	defer r.Body.Close()
	if r.StatusCode >= 400 {
        err = os.NewError(string(r.StatusCode))
        return
    }
    var root Root
    err = xml.Unmarshal(r.Body, &root)
	if err != nil {
		return
	}
	a := &root.Device
	if a.DeviceType != "urn:schemas-upnp-org:device:InternetGatewayDevice:1" {
	    err = os.NewError("No InternetGatewayDevice")
	    return
	}
	b := getChildDevice(a, "urn:schemas-upnp-org:device:WANDevice:1")
	if b == nil {
	    err = os.NewError("No WANDevice")
	    return
	}
	c := getChildDevice(b, "urn:schemas-upnp-org:device:WANConnectionDevice:1")
	if c == nil {
	    err = os.NewError("No WANConnectionDevice")
	    return
	}
	d := getChildService(c, "urn:schemas-upnp-org:service:WANIPConnection:1")
	if d == nil {
	    err = os.NewError("No WANIPConnection")
	    return
	}
	url = combineURL(rootURL, d.ControlURL)
	return
}

func combineURL(rootURL, subURL string) string {
    protocolEnd := "://"
    protoEndIndex := strings.Index(rootURL, protocolEnd)
    a := rootURL[protoEndIndex + len(protocolEnd):]
    rootIndex := strings.Index(a, "/")
    return rootURL[0:protoEndIndex + len(protocolEnd) + rootIndex] + subURL
}

type stringBuffer struct {
    base string
    current string
}

func NewStringBuffer(s string) *stringBuffer {
    return &stringBuffer{s, s}
}

func (sb *stringBuffer) Read(p []byte) (n int, err os.Error) {
    s := sb.current
    lenStr := len(s)
    if lenStr == 0 {
        return 0, os.EOF
    }
    n = len(p)
    if n > lenStr {
        n = lenStr
        err = os.EOF
    }
    for i := 0; i < n; i++ {
        p[i] = s[i]
    }
    sb.current = s[n:]
    return
}

func (sb *stringBuffer) Seek(offset int64, whence int) (ret int64, err os.Error) {
    var newOffset int64
    switch whence {
    case 0: // from beginning
        newOffset = offset
    case 1: // relative
        newOffset = int64(len(sb.base) - len(sb.current)) + offset
    case 2: // from end
        newOffset = int64(len(sb.base)) - offset
    default:
        err = os.NewError("bad whence")
        return
    }
    if newOffset < 0 || newOffset > int64(len(sb.base)) {
        err = os.NewError("offset out of range")
    } else {
        sb.current = sb.base[newOffset:]
        ret = newOffset
    }
    return
}

func soapRequest(url, function, message string) (r *http.Response, err os.Error) {
    fullMessage := "<?xml version=\"1.0\"?>" +
        "<s:Envelope xmlns:s=\"http://schemas.xmlsoap.org/soap/envelope/\" s:encodingStyle=\"http://schemas.xmlsoap.org/soap/encoding/\">\r\n" +
        "<s:Body>" + message + "</s:Body></s:Envelope>"

	var req http.Request
	req.Method = "POST"
	req.Body = NewStringBuffer(fullMessage)
	req.UserAgent = "Darwin/10.0.0, UPnP/1.0, MiniUPnPc/1.3"
	req.Header = map[string]string{
		"Content-Type": "text/xml",
		// "Transfer-Encoding": "chunked",
		"SOAPAction": "\"urn:schemas-upnp-org:service:WANIPConnection:1#" + function + "\"",
		"Connection": "Close",
		"Cache-Control": "no-cache",
		"Pragma": "no-cache",
	}

	req.URL, err = http.ParseURL(url)
	if err != nil {
		return
	}

	// log.Stderr("soapRequest ", req)

	r, err = http.Send(&req)

	if r.StatusCode >= 4000 {
        err = os.NewError("Error " + strconv.Itoa(r.StatusCode) + " for " + function)
        r.Body.Close()
        r = nil
        return
    }
	return
}

func (n *upnpNAT) ForwardPort(protocol string, externalPort, internalPort int,
    description string, timeout int) (err os.Error) {
    message := "<u:AddPortMapping xmlns:u=\"urn:schemas-upnp-org:service:WANIPConnection:1\">\r\n" +
        "<NewRemoteHost></NewRemoteHost><NewExternalPort>" + strconv.Itoa(externalPort) +
        "</NewExternalPort><NewProtocol>" + protocol + "</NewProtocol>" +
        "<NewInternalPort>" + strconv.Itoa(internalPort) + "</NewInternalPort><NewInternalClient>" +
        "192.168.0.124" + // TODO: Put our IP address here.
        "</NewInternalClient><NewEnabled>1</NewEnabled><NewPortMappingDescription>" +
        description +
        "</NewPortMappingDescription><NewLeaseDuration>" + strconv.Itoa(timeout) + "</NewLeaseDuration></u:AddPortMapping>"

    var response *http.Response
    response, err = soapRequest(n.serviceURL, "AddPortMapping", message)

    // TODO: check response to see if the port was forwarded
    _ = response
    return
}

func (n *upnpNAT) DeleteForwardingRule(protocol string, externalPort int) (err os.Error) {
    return
}

func testUPnP() {
    log.Stderr("Starting UPnP test")
    nat, err := Discover()
    log.Stderr("nat ", nat, "err ", err)
    if err != nil {
        return
    }
    port := 60001
    err = nat.ForwardPort("TCP", port, port, "Taipei-Torrent", 0)
    log.Stderr("err ", err)
    err = nat.DeleteForwardingRule("TCP", port)
    log.Stderr("err ", err)
}
<|endoftext|>


<commit_before>package tftest

import (
	"os"
	"path/filepath"
)

func symlinkFile(src string, dest string) (err error) {
	err = os.Symlink(src, dest)
	if err == nil {
		srcInfo, err := os.Stat(src)
		if err != nil {
			err = os.Chmod(dest, srcInfo.Mode())
		}
	}

	return
}

// symlinkDir is a simplistic function for recursively symlinking all files in a directory to a new path.
// It is intended only for limited internal use and does not cover all edge cases.
func symlinkDir(srcDir string, destDir string) (err error) {
	srcInfo, err := os.Stat(srcDir)
	if err != nil {
		return err
	}

	err = os.MkdirAll(destDir, srcInfo.Mode())
	if err != nil {
		return err
	}

	directory, _ := os.Open(srcDir)
	objects, err := directory.Readdir(-1)

	for _, obj := range objects {
		srcPath := filepath.Join(srcDir, obj.Name())
		destPath := filepath.Join(destDir, obj.Name())

		if obj.IsDir() {
			err = symlinkDir(srcPath, destPath)
			if err != nil {
				return err
			}
		} else {
			err = symlinkFile(srcPath, destPath)
			if err != nil {
				return err
			}
		}

	}
	return
}
<commit_msg>fix unclosed directory bug<commit_after>package tftest

import (
	"os"
	"path/filepath"
)

func symlinkFile(src string, dest string) (err error) {
	err = os.Symlink(src, dest)
	if err == nil {
		srcInfo, err := os.Stat(src)
		if err != nil {
			err = os.Chmod(dest, srcInfo.Mode())
		}
	}

	return
}

// symlinkDir is a simplistic function for recursively symlinking all files in a directory to a new path.
// It is intended only for limited internal use and does not cover all edge cases.
func symlinkDir(srcDir string, destDir string) (err error) {
	srcInfo, err := os.Stat(srcDir)
	if err != nil {
		return err
	}

	err = os.MkdirAll(destDir, srcInfo.Mode())
	if err != nil {
		return err
	}

	directory, _ := os.Open(srcDir)
	defer directory.Close()
	objects, err := directory.Readdir(-1)

	for _, obj := range objects {
		srcPath := filepath.Join(srcDir, obj.Name())
		destPath := filepath.Join(destDir, obj.Name())

		if obj.IsDir() {
			err = symlinkDir(srcPath, destPath)
			if err != nil {
				return err
			}
		} else {
			err = symlinkFile(srcPath, destPath)
			if err != nil {
				return err
			}
		}

	}
	return
}
<|endoftext|>


<commit_before>// Package snmputil provides helper routines for gosnmp

// Copyright 2016 Paul Stuart. All rights reserved.
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file.

package snmputil

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/soniah/gosnmp"
)

var (
	// Debug will log snmp debugging output if set
	Debug *log.Logger
	// Verbose will log output if set
	Verbose *log.Logger
	// ToOID is a lookup table to find the dotted form of a symbolic name
	ToOID = make(map[string]string)
	// Done will terminate all polling processes if closed
	Done      = make(chan struct{})
	oidLookup = make(map[string]int) // the index of the tuple
	table     = []oidTuple{}
	// how to break up column indexes with multiple elements
	multiName = strings.Fields("Grouping Member Element Item")
)

const (
	ifName  = ".1.3.6.1.2.1.31.1.1.1.1"
	ifAlias = ".1.3.6.1.2.1.31.1.1.1.18"
)

// Sender will send the interpreted PDU value to be saved or whathaveyou
type Sender func(string, map[string]string, interface{}, time.Time) error

// Criteria specifies what is to query and what to keep
type Criteria struct {
	OID     string            // OID can be dotted string or symbolic name
	Tags    map[string]string // any additional tags to associate
	Regexps []string          // filter resulting entries
	Keep    bool              // keep if resulting name matches, otherwise omit
}

// SnmpStats tracks SNMP request activity
type SnmpStats struct {
	LastError time.Time
	GetCnt    int64
	ErrCnt    int64
	Error     error
}

// StatsChan is used to request SnmpStats from a polling process
type StatsChan chan SnmpStats

// PDUFunc takes the original OID and a resulting SNMP packet to process
type PDUFunc func(string, gosnmp.SnmpPDU) error

// ErrFunc accepts SNMP errors for processing
type ErrFunc func(error)

// oidTuple is used to track the number of columns a table has
// It is used to determine if an OID is a table member
type oidTuple struct {
	OID, Name string
	Entries   int
}

// Tuples is a list of tuples containing OID table info
type Tuples []oidTuple

func (o Tuples) Len() int      { return len(o) }
func (o Tuples) Swap(i, j int) { o[i], o[j] = o[j], o[i] }
func (o Tuples) Less(i, j int) bool {
	return strings.Compare(o[i].OID, o[j].OID) < 1
}

func say(fmt string, args ...interface{}) {
	if Verbose != nil {
		Verbose.Printf(fmt, args...)
	}
}

// numerical returns the parsed data type in its numeric form
func numerical(s string) (interface{}, error) {
	if f, err := strconv.ParseFloat(s, 64); err == nil {
		return f, nil
	}
	if i, err := strconv.ParseInt(s, 0, 64); err == nil {
		return i, nil
	}
	return s, fmt.Errorf("not a number")
}

// LoadOIDs reads a file of OIDs and their symbolic names
func LoadOIDs(in io.Reader) error {
	scanner := bufio.NewScanner(in)
	for scanner.Scan() {
		f := strings.Fields(scanner.Text())
		if len(f) < 2 {
			continue
		}
		ToOID[f[0]] = f[1]
		table = append(table, oidTuple{f[1], f[0], 0})
	}
	// sort the tuples so we can find sequences of oids
	// with matching prefix
	sort.Sort(Tuples(table))
	for i := 0; i < len(table); i++ {
		t1 := table[i]
		cnt := 0
		for j := i + 1; j < len(table); j++ {
			t2 := table[j]
			if !strings.HasPrefix(t2.OID, t1.OID) {
				break
			}
			cnt++
		}
		table[i].Entries = cnt
	}
	for i, t := range table {
		oidLookup[t.OID] = i
	}
	return scanner.Err()
}

// LoadOIDFile is a helper routine to load OID descriptions from a file
func LoadOIDFile(filename string) error {
	f, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer f.Close()
	return LoadOIDs(f)
}

// makeString converts ascii octets into a string
func makeString(bits []string) string {
	chars := make([]byte, len(bits))
	for i, bit := range bits {
		n, _ := strconv.Atoi(bit)
		chars[i] = byte(n)
	}
	return string(chars)
}

// oidStrings converts ascii octets into an array of words
func oidStrings(in string) []string {
	words := []string{}
	bits := strings.Split(in, ".")
	for i := 0; i < len(bits); i++ {
		cnt, _ := strconv.Atoi(bits[i])
		end := i + cnt + 1
		if i > len(bits) || i >= end {
			break
		}
		if end > len(bits) {
			end = len(bits)
		}
		word := makeString(bits[i+1 : end])
		words = append(words, word)
		i += cnt
	}
	return words
}

// finder finds the longest matching table entry for base, the given OID
// TODO: this scheme should be replaced by a radix tree
func finder(base, oid string) (string, string, error) {
	if strings.HasPrefix(oid, ".") {
		oid = oid[1:]
	}
	i, ok := oidLookup[base]
	if !ok {
		return base, "(unknown)", fmt.Errorf("not found: %s", base)
	}
	tuple := table[i]
	for k := i + tuple.Entries; k > i; k-- {
		if strings.HasPrefix(oid, table[k].OID) {
			return table[k].OID, table[k].Name, nil
		}
	}
	return base, tuple.Name, nil
}

// Octets converts ascii octets into a byte array
func stringToOctets(in string) []byte {
	if strings.HasPrefix(in, ".") {
		in = in[1:]
	}
	bits := strings.Split(in, ".")
	reply := make([]byte, len(bits))
	for i, bit := range bits {
		b, _ := strconv.Atoi(bit)
		reply[i] = byte(b)
	}
	return reply
}

// Octets converts ascii octets into a byte array
func octetsToString(in []byte) string {
	buf := make([]string, len(in))
	for i, bit := range in {
		buf[i] = fmt.Sprintf("%d", bit)
	}
	return strings.Join(buf, ".")
}

// suffixValue returns a map the OID suffixes and their respective names for each column of table
func suffixValue(lookup map[string]string) PDUFunc {
	return func(root string, pdu gosnmp.SnmpPDU) error {
		switch pdu.Type {
		case gosnmp.OctetString:
			lookup[pdu.Name[len(root)+2:]] = string(pdu.Value.([]byte))
		case gosnmp.IPAddress:
			lookup[pdu.Name[len(root)+2:]] = pdu.Value.(string)
		default:
			say("UNKNOWN TYPE: %x VALUE: %v\n", pdu.Type, pdu.Value)
		}
		return nil
	}
}

// BulkColumns returns a WalkFunc that will process results from a bulkwalk
func BulkColumns(client *gosnmp.GoSNMP, crit Criteria, sender Sender) (gosnmp.WalkFunc, error) {
	filterNames := []*regexp.Regexp{}
	for _, n := range crit.Regexps {
		re, err := regexp.Compile(n)
		if err != nil {
			return nil, err
		}
		filterNames = append(filterNames, re)
	}

	columns := make(map[string]string)
	aliases := make(map[string]string)
	if err := BulkWalkAll(client, ifName, suffixValue(columns)); err != nil {
		return nil, err
	}
	if err := BulkWalkAll(client, ifAlias, suffixValue(aliases)); err != nil {
		return nil, err
	}

	return func(pdu gosnmp.SnmpPDU) error {
		// find the oid of a table entry, if it exists
		subOID, name, err := finder(crit.OID, pdu.Name)
		if err != nil {
			return err
		}

		filtered := crit.Keep
		for _, r := range filterNames {
			if r.MatchString(name) {
				if crit.Keep {
					filtered = false
					break
				}
				say("Omitting name: %s (%s)\n", name, subOID)
				return nil
			}
		}
		if filtered {
			say("Not keeping name: %s (%s)\n", name, subOID)
			return nil
		}

		var column, alias string
		suffix := pdu.Name[len(subOID)+2:]
		group := oidStrings(suffix)

		// interface names/aliases only apply to OIDs starting with 'if'
		if strings.HasPrefix(name, "if") {
			column = columns[suffix]
			alias = aliases[suffix]
		}
		if len(group) == 0 && len(column) == 0 && suffix != "0" {
			column = makeString(strings.Split(suffix, "."))
		}

		say("OID:%s SUFFIX:%s COL:%s NAME:%s TYPE:%x VALUE:%v\n", crit.OID, suffix, column, name, pdu.Type, pdu.Value)
		t := map[string]string{}
		if len(column) > 0 {
			t["Column"] = column
		}
		if len(alias) > 0 {
			t["Alias"] = alias
		}
		if len(group) > 0 && len(group[0]) > 0 {
			t["Grouping"] = group[0]
		}
		if len(group) > 1 && len(group[1]) > 0 {
			t["Member"] = group[1]
		}
		if len(group) > 3 && len(group[1]) > 0 {
			t["Element"] = group[2]
		}

		// copy tag values so we don't modify original tags map
		for k, v := range crit.Tags {
			t[k] = v
		}
		switch pdu.Type {
		case gosnmp.Integer, gosnmp.Counter32, gosnmp.Gauge32, gosnmp.TimeTicks, gosnmp.Counter64, gosnmp.Uinteger32:
		case gosnmp.IPAddress:
		case gosnmp.OctetString:
			s := string(pdu.Value.([]uint8))
			if n, err := numerical(s); err != nil {
				say("%s (%x) - non numerical: %s\n", name, pdu.Type, s)
				pdu.Value = n
			}
		default:
			say(name, "%s - unsupported type: %x value: %v\n", name, pdu.Type, pdu.Value)
			return nil
		}
		return sender(name, t, pdu.Value, time.Now())
	}, nil
}

// GetOID will return the OID representing name
func GetOID(oid string) (string, error) {
	if strings.HasPrefix(oid, ".") {
		oid = oid[1:]
	}
	if strings.HasPrefix(oid, "1.") {
		return oid, nil
	}
	fixed, ok := ToOID[oid]
	if !ok {
		return oid, fmt.Errorf("no OID found for %s", oid)
	}
	return fixed, nil
}

// BulkWalkAll applies bulk walk results to fn once all values returned (synchronously)
func BulkWalkAll(client *gosnmp.GoSNMP, oid string, fn PDUFunc) error {
	pdus, err := client.BulkWalkAll(oid)
	if err != nil {
		return err
	}
	for _, pdu := range pdus {
		if err := fn(oid, pdu); err != nil {
			return err
		}
	}
	return nil
}

// InterfaceNames will apply the interface name and oid to the given fn
func InterfaceNames(p Profile, fn func(string, string)) error {
	client, err := NewClient(p)
	if err != nil {
		return err
	}

	defer client.Conn.Close()
	return BulkWalkAll(client, ifName,
		func(root string, pdu gosnmp.SnmpPDU) error {
			switch pdu.Type {
			case gosnmp.OctetString:
				fn(pdu.Name, string(pdu.Value.([]byte)))
			}
			return nil
		})
}

// Bulkwalker will do a bulkwalk on the device specified in the Profile
func Bulkwalker(p Profile, crit Criteria, sender Sender, freq int, errFn ErrFunc, status chan StatsChan) error {
	client, err := NewClient(p)
	if err != nil {
		return err
	}
	crit.OID, err = GetOID(crit.OID)
	if err != nil {
		return err
	}
	if crit.Tags == nil {
		crit.Tags = make(map[string]string)
	}
	crit.Tags["Host"] = client.Target
	if Debug != nil {
		client.Logger = Debug
	}
	walker, err := BulkColumns(client, crit, sender)
	if err != nil {
		return err
	}
	go Poller(client, crit.OID, freq, walker, errFn, status)
	return nil
}

// Poller will make snmp requests indefinitely
func Poller(client *gosnmp.GoSNMP, oid string, freq int, walker gosnmp.WalkFunc, errFn ErrFunc, status chan StatsChan) {

	stats := SnmpStats{}

	c := time.Tick(time.Duration(freq) * time.Second)

	walk := func() {
		if err := client.BulkWalk(oid, walker); err != nil {
			if errFn != nil {
				errFn(err)
			}
			stats.ErrCnt++
			stats.LastError = time.Now()
		} else {
			stats.GetCnt++
		}
	}

	// poll immediately without waiting for first tick
	walk()

	for {
		select {
		case _ = <-c:
			walk()
		case s := <-status:
			s <- stats
		case _ = <-Done:
			client.Conn.Close()
			return
		}
	}
}
<commit_msg>use radix tree to fine OID table prefixes<commit_after>// Package snmputil provides helper routines for gosnmp

// Copyright 2016 Paul Stuart. All rights reserved.
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file.

package snmputil

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"

	radix "github.com/hashicorp/go-immutable-radix"
	"github.com/pkg/errors"
	"github.com/soniah/gosnmp"
)

var (
	// Debug will log snmp debugging output if set
	Debug *log.Logger
	// Verbose will log output if set
	Verbose *log.Logger
	// ToOID is a lookup table to find the dotted form of a symbolic name
	ToOID = make(map[string]string)
	// Done will terminate all polling processes if closed
	Done = make(chan struct{})
	// how to break up column indexes with multiple elements
	multiName = strings.Fields("Grouping Member Element Item")
	rtree     = radix.New()
)

const (
	ifName  = ".1.3.6.1.2.1.31.1.1.1.1"
	ifAlias = ".1.3.6.1.2.1.31.1.1.1.18"
)

// Sender will send the interpreted PDU value to be saved or whathaveyou
type Sender func(string, map[string]string, interface{}, time.Time) error

// Criteria specifies what is to query and what to keep
type Criteria struct {
	OID     string            // OID can be dotted string or symbolic name
	Tags    map[string]string // any additional tags to associate
	Regexps []string          // filter resulting entries
	Keep    bool              // keep if resulting name matches, otherwise omit
}

// SnmpStats tracks SNMP request activity
type SnmpStats struct {
	LastError time.Time
	GetCnt    int64
	ErrCnt    int64
	Error     error
}

// StatsChan is used to request SnmpStats from a polling process
type StatsChan chan SnmpStats

// ErrFunc accepts SNMP errors for processing
type ErrFunc func(error)

func say(fmt string, args ...interface{}) {
	if Verbose != nil {
		Verbose.Printf(fmt, args...)
	}
}

// numerical returns the parsed data type in its numeric form
func numerical(s string) (interface{}, error) {
	if f, err := strconv.ParseFloat(s, 64); err == nil {
		return f, nil
	}
	if i, err := strconv.ParseInt(s, 0, 64); err == nil {
		return i, nil
	}
	return s, fmt.Errorf("not a number")
}

// LoadOIDs reads a file of OIDs and their symbolic names
func LoadOIDs(in io.Reader) error {
	scanner := bufio.NewScanner(in)
	for scanner.Scan() {
		f := strings.Fields(scanner.Text())
		if len(f) < 2 {
			continue
		}
		// snmptranslate isn't providing leading dot
		if f[1][:1] != "." {
			f[1] = "." + f[1]
		}
		ToOID[f[0]] = f[1]
		rtree, _, _ = rtree.Insert([]byte(f[1]), f[0])
	}
	return scanner.Err()
}

// LoadOIDFile is a helper routine to load OID descriptions from a file
func LoadOIDFile(filename string) error {
	f, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer f.Close()
	return LoadOIDs(f)
}

// makeString converts ascii octets into a string
func makeString(bits []string) string {
	chars := make([]byte, len(bits))
	for i, bit := range bits {
		n, _ := strconv.Atoi(bit)
		chars[i] = byte(n)
	}
	return string(chars)
}

// oidStrings converts ascii octets into an array of words
func oidStrings(in string) []string {
	words := []string{}
	bits := strings.Split(in, ".")
	for i := 0; i < len(bits); i++ {
		cnt, _ := strconv.Atoi(bits[i])
		end := i + cnt + 1
		if i > len(bits) || i >= end {
			break
		}
		if end > len(bits) {
			end = len(bits)
		}
		word := makeString(bits[i+1 : end])
		words = append(words, word)
		i += cnt
	}
	return words
}

// Octets converts ascii octets into a byte array
func stringToOctets(in string) []byte {
	if strings.HasPrefix(in, ".") {
		in = in[1:]
	}
	bits := strings.Split(in, ".")
	reply := make([]byte, len(bits))
	for i, bit := range bits {
		b, _ := strconv.Atoi(bit)
		reply[i] = byte(b)
	}
	return reply
}

// Octets converts ascii octets into a byte array
func octetsToString(in []byte) string {
	buf := make([]string, len(in))
	for i, bit := range in {
		buf[i] = fmt.Sprintf("%d", bit)
	}
	return strings.Join(buf, ".")
}

// BulkColumns returns a WalkFunc that will process results from a bulkwalk
func BulkColumns(client *gosnmp.GoSNMP, crit Criteria, sender Sender) (gosnmp.WalkFunc, error) {
	filterNames := []*regexp.Regexp{}
	for _, n := range crit.Regexps {
		re, err := regexp.Compile(n)
		if err != nil {
			return nil, err
		}
		filterNames = append(filterNames, re)
	}

	// suffixValue returns a map the OID indexes and their respective names for each column of table
	suffixValue := func(root string, lookup map[string]string) (gosnmp.WalkFunc, error) {
		oid, err := GetOID(root)
		if err != nil {
			return nil, err
		}
		return func(pdu gosnmp.SnmpPDU) error {
			switch pdu.Type {
			case gosnmp.OctetString:
				lookup[pdu.Name[len(oid)+2:]] = string(pdu.Value.([]byte))
			case gosnmp.IPAddress:
				lookup[pdu.Name[len(oid)+2:]] = pdu.Value.(string)
			default:
				say("UNKNOWN TYPE: %x VALUE: %v\n", pdu.Type, pdu.Value)
			}
			return nil
		}, nil
	}
	columns := make(map[string]string)
	suffixFn, err := suffixValue(ifName, columns)
	if err != nil {
		return nil, err
	}
	if err := BulkWalkAll(client, ifName, suffixFn); err != nil {
		return nil, err
	}

	aliases := make(map[string]string)
	suffixFn, err = suffixValue(ifAlias, aliases)
	if err != nil {
		return nil, err
	}
	if err := BulkWalkAll(client, ifAlias, suffixFn); err != nil {
		return nil, err
	}

	return func(pdu gosnmp.SnmpPDU) error {
		subOID, v, ok := rtree.Root().LongestPrefix([]byte(pdu.Name))
		if !ok {
			return errors.Errorf("cannot find name for OID: %s", pdu.Name)
		}
		name := v.(string)

		filtered := crit.Keep
		for _, r := range filterNames {
			if r.MatchString(name) {
				if crit.Keep {
					filtered = false
					break
				}
				say("Omitting name: %s (%s)\n", name, subOID)
				return nil
			}
		}
		if filtered {
			say("Not keeping name: %s (%s)\n", name, subOID)
			return nil
		}

		var column, alias string
		suffix := pdu.Name[len(subOID)+1:]
		group := oidStrings(suffix)

		// interface names/aliases only apply to OIDs starting with 'if'
		if strings.HasPrefix(name, "if") {
			column = columns[suffix]
			alias = aliases[suffix]
		}
		if len(group) == 0 && len(column) == 0 && suffix != "0" {
			column = makeString(strings.Split(suffix, "."))
		}

		say("OID:%s SUFFIX:%s COL:%s NAME:%s TYPE:%x VALUE:%v\n", crit.OID, suffix, column, name, pdu.Type, pdu.Value)
		t := map[string]string{}
		if len(column) > 0 {
			t["Column"] = column
		}
		if len(alias) > 0 {
			t["Alias"] = alias
		}
		if len(group) > 0 && len(group[0]) > 0 {
			t["Grouping"] = group[0]
		}
		if len(group) > 1 && len(group[1]) > 0 {
			t["Member"] = group[1]
		}
		if len(group) > 3 && len(group[1]) > 0 {
			t["Element"] = group[2]
		}

		// copy tag values so we don't modify original tags map
		for k, v := range crit.Tags {
			t[k] = v
		}
		switch pdu.Type {
		case gosnmp.Integer, gosnmp.Counter32, gosnmp.Gauge32, gosnmp.TimeTicks, gosnmp.Counter64, gosnmp.Uinteger32:
		case gosnmp.IPAddress:
		case gosnmp.OctetString:
			s := string(pdu.Value.([]uint8))
			if n, err := numerical(s); err != nil {
				say("%s (%x) - non numerical: %s\n", name, pdu.Type, s)
				pdu.Value = n
			}
		default:
			say(name, "%s - unsupported type: %x value: %v\n", name, pdu.Type, pdu.Value)
			return nil
		}
		return sender(name, t, pdu.Value, time.Now())
	}, nil
}

// GetOID will return the OID representing name
func GetOID(oid string) (string, error) {
	if strings.HasPrefix(oid, ".") {
		oid = oid[1:]
	}
	if strings.HasPrefix(oid, "1.") {
		return oid, nil
	}
	fixed, ok := ToOID[oid]
	if !ok {
		return oid, fmt.Errorf("no OID found for %s", oid)
	}
	return fixed, nil
}

// BulkWalkAll applies bulk walk results to fn once all values returned (synchronously)
func BulkWalkAll(client *gosnmp.GoSNMP, oid string, fn gosnmp.WalkFunc) error {
	pdus, err := client.BulkWalkAll(oid)
	if err != nil {
		return err
	}
	for _, pdu := range pdus {
		if err := fn(pdu); err != nil {
			return err
		}
	}
	return nil
}

// InterfaceNames will apply the interface name and oid to the given fn
func InterfaceNames(p Profile, fn func(string, string)) error {
	client, err := NewClient(p)
	if err != nil {
		return err
	}

	defer client.Conn.Close()
	return BulkWalkAll(client, ifName,
		func(pdu gosnmp.SnmpPDU) error {
			switch pdu.Type {
			case gosnmp.OctetString:
				fn(pdu.Name, string(pdu.Value.([]byte)))
			}
			return nil
		})
}

// Bulkwalker will do a bulkwalk on the device specified in the Profile
func Bulkwalker(p Profile, crit Criteria, sender Sender, freq int, errFn ErrFunc, status chan StatsChan) error {
	client, err := NewClient(p)
	if err != nil {
		return err
	}
	crit.OID, err = GetOID(crit.OID)
	if err != nil {
		return err
	}
	if crit.Tags == nil {
		crit.Tags = make(map[string]string)
	}
	crit.Tags["Host"] = client.Target
	if Debug != nil {
		client.Logger = Debug
	}
	walker, err := BulkColumns(client, crit, sender)
	if err != nil {
		return err
	}
	go Poller(client, crit.OID, freq, walker, errFn, status)
	return nil
}

// Poller will make snmp requests indefinitely
func Poller(client *gosnmp.GoSNMP, oid string, freq int, walker gosnmp.WalkFunc, errFn ErrFunc, status chan StatsChan) {

	stats := SnmpStats{}

	c := time.Tick(time.Duration(freq) * time.Second)

	walk := func() {
		if err := client.BulkWalk(oid, walker); err != nil {
			if errFn != nil {
				errFn(err)
			}
			stats.ErrCnt++
			stats.LastError = time.Now()
		} else {
			stats.GetCnt++
		}
	}

	// poll immediately without waiting for first tick
	walk()

	for {
		select {
		case _ = <-c:
			walk()
		case s := <-status:
			s <- stats
		case _ = <-Done:
			client.Conn.Close()
			return
		}
	}
}
<|endoftext|>


<commit_before>package uuid

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

// GenerateUUID is used to generate a random UUID
func GenerateUUID() (string, error) {
	buf := make([]byte, 16)
	if _, err := rand.Read(buf); err != nil {
		return "", fmt.Errorf("failed to read random bytes: %v", err)
	}

	return FormatUUID(buf)
}

func FormatUUID(buf []byte) (string, error) {
	if len(buf) != 16 {
		return "", fmt.Errorf("wrong length byte slice (%d)", len(buf))
	}

	return fmt.Sprintf("%08x-%04x-%04x-%04x-%12x",
		buf[0:4],
		buf[4:6],
		buf[6:8],
		buf[8:10],
		buf[10:16]), nil
}

func ParseUUID(uuid string) ([]byte, error) {
	if len(uuid) != 36 {
		return nil, fmt.Errorf("uuid string is wrong length")
	}

	hyph := []byte("-")

	if uuid[8] != hyph[0] ||
		uuid[13] != hyph[0] ||
		uuid[18] != hyph[0] ||
		uuid[23] != hyph[0] {
		return nil, fmt.Errorf("uuid is improperly formatted")
	}

	hexStr := uuid[0:8] + uuid[9:13] + uuid[14:18] + uuid[19:23] + uuid[24:36]

	ret, err := hex.DecodeString(hexStr)
	if err != nil {
		return nil, err
	}
	if len(ret) != 16 {
		return nil, fmt.Errorf("decoded hex is the wrong length")
	}

	return ret, nil
}
<commit_msg>Added GenerateRandomBytes(int)<commit_after>package uuid

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

// GenerateRandomBytes is used to generate random bytes of given size.
func GenerateRandomBytes(size int) ([]byte, error) {
	buf := make([]byte, size)
	if _, err := rand.Read(buf); err != nil {
		return nil, fmt.Errorf("failed to read random bytes: %v", err)
	}
	return buf, nil
}

// GenerateUUID is used to generate a random UUID
func GenerateUUID() (string, error) {
	buf, err := GenerateRandomBytes(16)
	if err != nil {
		return "", err
	}
	return FormatUUID(buf)
}

func FormatUUID(buf []byte) (string, error) {
	if len(buf) != 16 {
		return "", fmt.Errorf("wrong length byte slice (%d)", len(buf))
	}

	return fmt.Sprintf("%08x-%04x-%04x-%04x-%12x",
		buf[0:4],
		buf[4:6],
		buf[6:8],
		buf[8:10],
		buf[10:16]), nil
}

func ParseUUID(uuid string) ([]byte, error) {
	if len(uuid) != 36 {
		return nil, fmt.Errorf("uuid string is wrong length")
	}

	hyph := []byte("-")

	if uuid[8] != hyph[0] ||
		uuid[13] != hyph[0] ||
		uuid[18] != hyph[0] ||
		uuid[23] != hyph[0] {
		return nil, fmt.Errorf("uuid is improperly formatted")
	}

	hexStr := uuid[0:8] + uuid[9:13] + uuid[14:18] + uuid[19:23] + uuid[24:36]

	ret, err := hex.DecodeString(hexStr)
	if err != nil {
		return nil, err
	}
	if len(ret) != 16 {
		return nil, fmt.Errorf("decoded hex is the wrong length")
	}

	return ret, nil
}
<|endoftext|>


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

package state

import (
	"reflect"

	"github.com/juju/utils/set"
	gc "gopkg.in/check.v1"
)

type MigrationSuite struct{}

var _ = gc.Suite(&MigrationSuite{})

func (s *MigrationSuite) TestKnownCollections(c *gc.C) {
	completedCollections := set.NewStrings(
		modelsC,
		modelUsersC,
		modelUserLastConnectionC,
		settingsC,
		statusesC,

		// machine
		instanceDataC,
		machinesC,

		// service
		servicesC,
		// settings reference counts are only used for services
		settingsrefsC,
	)

	ignoredCollections := set.NewStrings(
		// We don't export the controller model at this stage.
		controllersC,
		// Users aren't migrated.
		usersC,
		userLastLoginC,
		// userenvnameC is just to provide a unique key constraint.
		usermodelnameC,
		// Metrics aren't migrated.
		metricsC,
		// leaseC is deprecated in favour of leasesC.
		leaseC,
		// Backup and restore information is not migrated.
		restoreInfoC,
		// upgradeInfoC is used to coordinate upgrades and schema migrations,
		// and aren't needed for model migrations.
		upgradeInfoC,
		// Not exported, but the tools will possibly need to be either bundled
		// with the representation or sent separately.
		toolsmetadataC,
		// Transaction stuff.
		"txns",
		"txns.log",

		// The container ref document is primarily there to keep track
		// of a particular machine's containers. The migration format
		// uses object containment for this purpose.
		containerRefsC,
		// The min units collection is only used to trigger a watcher
		// in order to have the service add or remove units if the minimum
		// number of units is changed. The Service doc has all we need
		// for migratino.
		minUnitsC,
	)

	// THIS SET WILL BE REMOVED WHEN MIGRATIONS ARE COMPLETE
	todoCollections := set.NewStrings(
		// model
		blocksC,
		cleanupsC,
		cloudimagemetadataC,
		sequenceC,

		// machine
		rebootC,

		// service / unit
		assignUnitC,
		charmsC,
		leasesC,
		openedPortsC,
		"payloads",
		unitsC,

		// relation
		relationsC,
		relationScopesC,

		// storage
		blockDevicesC,
		filesystemsC,
		filesystemAttachmentsC,
		storageInstancesC,
		storageAttachmentsC,
		storageConstraintsC,
		volumesC,
		volumeAttachmentsC,

		// network
		ipaddressesC,
		networksC,
		networkInterfacesC,
		requestedNetworksC,
		subnetsC,
		spacesC,

		// actions
		actionsC,
		actionNotificationsC,
		actionresultsC,

		// done as part of machines/services/units
		annotationsC,
		constraintsC,
		statusesHistoryC,

		// uncategorised
		meterStatusC,    // red / green status for metrics / charms
		metricsManagerC, // should really be copied across

	)

	envCollections := set.NewStrings()
	for name := range allCollections() {
		envCollections.Add(name)
	}

	known := completedCollections.Union(ignoredCollections)

	remainder := envCollections.Difference(known)
	remainder = remainder.Difference(todoCollections)

	// If this test fails, it means that a new collection has been added
	// but migrations for it has not been done. This is a Bad Thing™.
	c.Assert(remainder, gc.HasLen, 0)
}

func (s *MigrationSuite) TestModelDocFields(c *gc.C) {
	fields := set.NewStrings(
		// UUID and Mame are constructed from the model config.
		"UUID",
		"Name",
		// Life will always be alive, or we won't be migrating.
		"Life",
		"Owner",
		"LatestAvailableTools",
		// ServerUUID is recreated when the new model is created in the
		// new controller (yay name changes).
		"ServerUUID",
		// Both of the times for dying and death are empty as the model
		// is alive.
		"TimeOfDying",
		"TimeOfDeath",
	)
	s.AssertExportedFields(c, modelDoc{}, fields)
}

func (s *MigrationSuite) TestEnvUserDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ID is the same as UserName (but lowercased)
		"ID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Tracked fields:
		"UserName",
		"DisplayName",
		"CreatedBy",
		"DateCreated",
		"ReadOnly",
	)
	s.AssertExportedFields(c, modelUserDoc{}, fields)
}

func (s *MigrationSuite) TestEnvUserLastConnectionDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ID is the same as UserName (but lowercased)
		"ID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// UserName is captured in the migration.User.
		"UserName",
		"LastConnection",
	)
	s.AssertExportedFields(c, modelUserLastConnectionDoc{}, fields)
}

func (s *MigrationSuite) TestMachineDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + machine id
		"DocID",
		// ID is the machine id
		"Id",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Life is always alive, confirmed by export precheck.
		"Life",

		"Addresses",
		"ContainerType",
		"Jobs",
		"MachineAddresses",
		"Nonce",
		"PasswordHash",
		"Placement",
		"PreferredPrivateAddress",
		"PreferredPublicAddress",
		"Series",
		"SupportedContainers",
		"SupportedContainersKnown",
		"Tools",
	)
	todo := set.NewStrings(
		"Principals",
		"Volumes",
		"NoVote",
		"Clean",
		"Filesystems",
		"HasVote",
	)
	s.AssertExportedFields(c, machineDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestInstanceDataFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + machine id
		"DocID",
		"MachineId",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",

		"InstanceId",
		"Status",
		"Arch",
		"Mem",
		"RootDisk",
		"CpuCores",
		"CpuPower",
		"Tags",
		"AvailZone",
	)
	s.AssertExportedFields(c, instanceData{}, fields)
}

func (s *MigrationSuite) TestServiceDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + name
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Always alive, not explicitly exported.
		"Life",
		// OwnerTag is deprecated and should be deleted.
		"OwnerTag",
		// TxnRevno is mgo internals and should not be migrated.
		"TxnRevno",

		"Name",
		"Series",
		"Subordinate",
		"CharmURL",
		"ForceCharm",
		"Exposed",
		"MinUnits",
	)
	todo := set.NewStrings(
		"UnitCount",
		"RelationCount",
		"MetricCredentials",
	)
	s.AssertExportedFields(c, serviceDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestSettingsRefsDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",

		"RefCount",
	)
	s.AssertExportedFields(c, settingsRefsDoc{}, fields)
}

func (s *MigrationSuite) AssertExportedFields(c *gc.C, doc interface{}, fields set.Strings) {
	expected := getExportedFields(doc)
	unknown := expected.Difference(fields)
	// If this test fails, it means that extra fields have been added to the
	// doc without thinking about the migration implications.
	c.Assert(unknown, gc.HasLen, 0)
}

func getExportedFields(arg interface{}) set.Strings {
	t := reflect.TypeOf(arg)
	result := set.NewStrings()

	count := t.NumField()
	for i := 0; i < count; i++ {
		f := t.Field(i)
		// empty PkgPath means exported field.
		// see https://golang.org/pkg/reflect/#StructField
		if f.PkgPath == "" {
			result.Add(f.Name)
		}
	}

	return result
}
<commit_msg>Test for unitDoc.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state

import (
	"reflect"

	"github.com/juju/utils/set"
	gc "gopkg.in/check.v1"
)

type MigrationSuite struct{}

var _ = gc.Suite(&MigrationSuite{})

func (s *MigrationSuite) TestKnownCollections(c *gc.C) {
	completedCollections := set.NewStrings(
		modelsC,
		modelUsersC,
		modelUserLastConnectionC,
		settingsC,
		statusesC,

		// machine
		instanceDataC,
		machinesC,

		// service / unit
		servicesC,
		unitsC,

		// settings reference counts are only used for services
		settingsrefsC,
	)

	ignoredCollections := set.NewStrings(
		// We don't export the controller model at this stage.
		controllersC,
		// Users aren't migrated.
		usersC,
		userLastLoginC,
		// userenvnameC is just to provide a unique key constraint.
		usermodelnameC,
		// Metrics aren't migrated.
		metricsC,
		// leaseC is deprecated in favour of leasesC.
		leaseC,
		// Backup and restore information is not migrated.
		restoreInfoC,
		// upgradeInfoC is used to coordinate upgrades and schema migrations,
		// and aren't needed for model migrations.
		upgradeInfoC,
		// Not exported, but the tools will possibly need to be either bundled
		// with the representation or sent separately.
		toolsmetadataC,
		// Transaction stuff.
		"txns",
		"txns.log",

		// The container ref document is primarily there to keep track
		// of a particular machine's containers. The migration format
		// uses object containment for this purpose.
		containerRefsC,
		// The min units collection is only used to trigger a watcher
		// in order to have the service add or remove units if the minimum
		// number of units is changed. The Service doc has all we need
		// for migratino.
		minUnitsC,
	)

	// THIS SET WILL BE REMOVED WHEN MIGRATIONS ARE COMPLETE
	todoCollections := set.NewStrings(
		// model
		blocksC,
		cleanupsC,
		cloudimagemetadataC,
		sequenceC,

		// machine
		rebootC,

		// service / unit
		assignUnitC,
		charmsC,
		leasesC,
		openedPortsC,
		"payloads",

		// relation
		relationsC,
		relationScopesC,

		// storage
		blockDevicesC,
		filesystemsC,
		filesystemAttachmentsC,
		storageInstancesC,
		storageAttachmentsC,
		storageConstraintsC,
		volumesC,
		volumeAttachmentsC,

		// network
		ipaddressesC,
		networksC,
		networkInterfacesC,
		requestedNetworksC,
		subnetsC,
		spacesC,

		// actions
		actionsC,
		actionNotificationsC,
		actionresultsC,

		// done as part of machines/services/units
		annotationsC,
		constraintsC,
		statusesHistoryC,

		// uncategorised
		meterStatusC,    // red / green status for metrics / charms
		metricsManagerC, // should really be copied across

	)

	envCollections := set.NewStrings()
	for name := range allCollections() {
		envCollections.Add(name)
	}

	known := completedCollections.Union(ignoredCollections)

	remainder := envCollections.Difference(known)
	remainder = remainder.Difference(todoCollections)

	// If this test fails, it means that a new collection has been added
	// but migrations for it has not been done. This is a Bad Thing™.
	c.Assert(remainder, gc.HasLen, 0)
}

func (s *MigrationSuite) TestModelDocFields(c *gc.C) {
	fields := set.NewStrings(
		// UUID and Mame are constructed from the model config.
		"UUID",
		"Name",
		// Life will always be alive, or we won't be migrating.
		"Life",
		"Owner",
		"LatestAvailableTools",
		// ServerUUID is recreated when the new model is created in the
		// new controller (yay name changes).
		"ServerUUID",
		// Both of the times for dying and death are empty as the model
		// is alive.
		"TimeOfDying",
		"TimeOfDeath",
	)
	s.AssertExportedFields(c, modelDoc{}, fields)
}

func (s *MigrationSuite) TestEnvUserDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ID is the same as UserName (but lowercased)
		"ID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Tracked fields:
		"UserName",
		"DisplayName",
		"CreatedBy",
		"DateCreated",
		"ReadOnly",
	)
	s.AssertExportedFields(c, modelUserDoc{}, fields)
}

func (s *MigrationSuite) TestEnvUserLastConnectionDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ID is the same as UserName (but lowercased)
		"ID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// UserName is captured in the migration.User.
		"UserName",
		"LastConnection",
	)
	s.AssertExportedFields(c, modelUserLastConnectionDoc{}, fields)
}

func (s *MigrationSuite) TestMachineDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + machine id
		"DocID",
		// ID is the machine id
		"Id",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Life is always alive, confirmed by export precheck.
		"Life",

		"Addresses",
		"ContainerType",
		"Jobs",
		"MachineAddresses",
		"Nonce",
		"PasswordHash",
		"Placement",
		"PreferredPrivateAddress",
		"PreferredPublicAddress",
		"Series",
		"SupportedContainers",
		"SupportedContainersKnown",
		"Tools",
	)
	todo := set.NewStrings(
		"Principals",
		"Volumes",
		"NoVote",
		"Clean",
		"Filesystems",
		"HasVote",
	)
	s.AssertExportedFields(c, machineDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestInstanceDataFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + machine id
		"DocID",
		"MachineId",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",

		"InstanceId",
		"Status",
		"Arch",
		"Mem",
		"RootDisk",
		"CpuCores",
		"CpuPower",
		"Tags",
		"AvailZone",
	)
	s.AssertExportedFields(c, instanceData{}, fields)
}

func (s *MigrationSuite) TestServiceDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID is the env + name
		"DocID",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Always alive, not explicitly exported.
		"Life",
		// OwnerTag is deprecated and should be deleted.
		"OwnerTag",
		// TxnRevno is mgo internals and should not be migrated.
		"TxnRevno",

		"Name",
		"Series",
		"Subordinate",
		"CharmURL",
		"ForceCharm",
		"Exposed",
		"MinUnits",
	)
	todo := set.NewStrings(
		"UnitCount",
		"RelationCount",
		"MetricCredentials",
	)
	s.AssertExportedFields(c, serviceDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) TestSettingsRefsDocFields(c *gc.C) {
	fields := set.NewStrings(
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",

		"RefCount",
	)
	s.AssertExportedFields(c, settingsRefsDoc{}, fields)
}

func (s *MigrationSuite) TestUnitDocFields(c *gc.C) {
	fields := set.NewStrings(
		// DocID itself isn't migrated
		"DocID",
		"Name",
		// ModelUUID shouldn't be exported, and is inherited
		// from the model definition.
		"ModelUUID",
		// Service is implicit in the migration structure through containment.
		"Service",
		// Series and CharmURL also come from the service.
		"Series",
		"CharmURL",
		"Principal",
		"Subordinates",
		"MachineId",
		// Resolved is not migrated as we check that all is good before we start.
		"Resolved",
		"Tools",
		// Life isn't migrated as we only migrate live things.
		"Life",
		// TxnRevno isn't migrated.
		"TxnRevno",
		"PasswordHash",
		// Obsolete and not migrated.
		"Ports",
		"PublicAddress",
		"PrivateAddress",
	)
	todo := set.NewStrings(
		"StorageAttachmentCount",
	)

	s.AssertExportedFields(c, unitDoc{}, fields.Union(todo))
}

func (s *MigrationSuite) AssertExportedFields(c *gc.C, doc interface{}, fields set.Strings) {
	expected := getExportedFields(doc)
	unknown := expected.Difference(fields)
	// If this test fails, it means that extra fields have been added to the
	// doc without thinking about the migration implications.
	c.Assert(unknown, gc.HasLen, 0)
}

func getExportedFields(arg interface{}) set.Strings {
	t := reflect.TypeOf(arg)
	result := set.NewStrings()

	count := t.NumField()
	for i := 0; i < count; i++ {
		f := t.Field(i)
		// empty PkgPath means exported field.
		// see https://golang.org/pkg/reflect/#StructField
		if f.PkgPath == "" {
			result.Add(f.Name)
		}
	}

	return result
}
<|endoftext|>


<commit_before>package trafo

import (
	"fmt"
	. "github.com/101loops/bdd"

	ds "appengine/datastore"
)

type loadEntity struct {
	A string
	B int

	beforeFunc func() error
	afterFunc  func() error
}

func (l *loadEntity) BeforeLoad() error {
	if l.beforeFunc != nil {
		return l.beforeFunc()
	}
	return nil
}

func (l *loadEntity) AfterLoad() error {
	if l.afterFunc != nil {
		return l.afterFunc()
	}
	return nil
}

var _ = Describe("Doc: Load", func() {

	var entity loadEntity

	BeforeEach(func() {
		entity = loadEntity{}
		CodecSet.AddMust(entity)
	})

	It("should load an entity from channel properties", func() {
		doc, err := newDocFromInst(&entity)
		Check(err, IsNil)
		Check(doc, NotNil)

		c := newPropertyChannel()
		err = doc.Load(c)
		Check(err, IsNil)
		Check(c, IsClosed)

		res := (doc.get()).(*loadEntity)
		Check(res.A, Equals, "abc")
		Check(res.B, EqualsNum, 1)
	})

	It("should return an error when BeforeLoad errs", func() {
		entity.beforeFunc = func() error {
			return fmt.Errorf("an error")
		}

		doc, err := newDocFromInst(&entity)
		Check(err, IsNil)
		Check(doc, NotNil)

		c := newPropertyChannel()
		err = doc.Load(c)
		Check(err, ErrorContains, "an error")
		Check(c, IsClosed)

		res := (doc.get()).(*loadEntity)
		Check(res.A, Equals, "")
	})

	It("should return an error when AfterLoad errs", func() {
		entity.afterFunc = func() error {
			return fmt.Errorf("an error")
		}

		doc, err := newDocFromInst(&entity)
		Check(err, IsNil)
		Check(doc, NotNil)

		c := newPropertyChannel()
		err = doc.Load(c)
		Check(err, ErrorContains, "an error")
		Check(c, IsClosed)

		res := (doc.get()).(*loadEntity)
		Check(res.A, Equals, "abc")
		Check(res.B, EqualsNum, 1)
	})
})

func newPropertyChannel() chan ds.Property {
	c := make(chan ds.Property)
	go func() {
		c <- ds.Property{Name: "A", Value: "abc"}
		c <- ds.Property{Name: "B", Value: int64(1)}
		close(c)
	}()
	return c
}
<commit_msg>test loading error<commit_after>package trafo

import (
	"fmt"
	. "github.com/101loops/bdd"

	ds "appengine/datastore"
)

type loadEntity struct {
	A string
	B int

	beforeFunc func() error
	afterFunc  func() error
}

func (l *loadEntity) BeforeLoad() error {
	if l.beforeFunc != nil {
		return l.beforeFunc()
	}
	return nil
}

func (l *loadEntity) AfterLoad() error {
	if l.afterFunc != nil {
		return l.afterFunc()
	}
	return nil
}

var _ = Describe("Doc: Load", func() {

	var entity loadEntity
	var validProps = []ds.Property{
		{Name: "A", Value: "abc"},
		{Name: "B", Value: int64(1)},
	}

	BeforeEach(func() {
		entity = loadEntity{}
		CodecSet.AddMust(entity)
	})

	It("should load an entity from channel properties", func() {
		doc, err := newDocFromInst(&entity)
		Check(err, IsNil)
		Check(doc, NotNil)

		c := newPropertyChannel(validProps)
		err = doc.Load(c)
		Check(err, IsNil)
		Check(c, IsClosed)

		res := (doc.get()).(*loadEntity)
		Check(res.A, Equals, "abc")
		Check(res.B, EqualsNum, 1)
	})

	It("should return an error when loading fails", func() {
		doc, _ := newDocFromInst(&entity)

		c := newPropertyChannel([]ds.Property{{Name: "Invalid"}})
		err := doc.Load(c)

		Check(err, ErrorContains, `cannot load field "Invalid" into a "trafo.loadEntity": no such struct field`)
		Check(c, IsClosed)
	})

	It("should return an error when BeforeLoad fails", func() {
		entity.beforeFunc = func() error {
			return fmt.Errorf("an error")
		}

		doc, _ := newDocFromInst(&entity)
		c := newPropertyChannel(validProps)

		err := doc.Load(c)
		Check(err, ErrorContains, "an error")
		Check(c, IsClosed)

		res := (doc.get()).(*loadEntity)
		Check(res.A, Equals, "")
	})

	It("should return an error when AfterLoad fails", func() {
		entity.afterFunc = func() error {
			return fmt.Errorf("an error")
		}

		doc, _ := newDocFromInst(&entity)
		c := newPropertyChannel(validProps)

		err := doc.Load(c)
		Check(err, ErrorContains, "an error")
		Check(c, IsClosed)

		res := (doc.get()).(*loadEntity)
		Check(res.A, Equals, "abc")
		Check(res.B, EqualsNum, 1)
	})
})

func newPropertyChannel(props []ds.Property) chan ds.Property {
	c := make(chan ds.Property)
	go func() {
		for _, p := range props {
			c <- p
		}
		close(c)
	}()
	return c
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	mapset "github.com/deckarep/golang-set"
)

func main() {

	s1 := mapset.NewSetFromSlice([]interface{}{"Delhi", "Gurgaon", "Noida", "Greater Noida", "Ghaziabad", "Manesar", "Faridabad",
		"Sonipat", "Meerut", "Chandigarh", "Ambala", "Bharatpur", "Mathura", "Rishikesh", "Dehradun",
		"Haridwar", "Mussoorie"})

	fmt.Println(s1.Contains("Delhi"))
	fmt.Println(s1.Contains("XYZ"))

}
<commit_msg>updated sets example<commit_after>package main

import (
	"fmt"
	mapset "github.com/deckarep/golang-set"
)

func main() {
	// initialiazing with some data. can also initialize with empty
	s1 := mapset.NewSetFromSlice([]interface{}{"Delhi", "Gurgaon"})

	// adding more values
	s1.Add("Noida")
	// adding duplicate values
	s1.Add("Noida")

	// convert to slice
	fmt.Println(s1.ToSlice()) // duplicates will be removed

	// check for existance of a value
	fmt.Println(s1.Contains("Delhi"))
	fmt.Println(s1.Contains("XYZ"))
}
<|endoftext|>


<commit_before>// +build linux

package systemd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	systemd "github.com/coreos/go-systemd/dbus"
	"github.com/docker/libcontainer/cgroups"
	"github.com/docker/libcontainer/cgroups/fs"
	"github.com/docker/libcontainer/configs"
	"github.com/godbus/dbus"
)

type Manager struct {
	Cgroups *configs.Cgroup
	Paths   map[string]string
}

type subsystem interface {
	// Returns the stats, as 'stats', corresponding to the cgroup under 'path'.
	GetStats(path string, stats *cgroups.Stats) error
	// Set the cgroup represented by cgroup.
	Set(path string, cgroup *configs.Cgroup) error
}

var subsystems = map[string]subsystem{
	"devices":    &fs.DevicesGroup{},
	"memory":     &fs.MemoryGroup{},
	"cpu":        &fs.CpuGroup{},
	"cpuset":     &fs.CpusetGroup{},
	"cpuacct":    &fs.CpuacctGroup{},
	"blkio":      &fs.BlkioGroup{},
	"perf_event": &fs.PerfEventGroup{},
	"freezer":    &fs.FreezerGroup{},
}

const (
	testScopeWait = 4
)

var (
	connLock                        sync.Mutex
	theConn                         *systemd.Conn
	hasStartTransientUnit           bool
	hasTransientDefaultDependencies bool
)

func newProp(name string, units interface{}) systemd.Property {
	return systemd.Property{
		Name:  name,
		Value: dbus.MakeVariant(units),
	}
}

func UseSystemd() bool {
	s, err := os.Stat("/run/systemd/system")
	if err != nil || !s.IsDir() {
		return false
	}

	connLock.Lock()
	defer connLock.Unlock()

	if theConn == nil {
		var err error
		theConn, err = systemd.New()
		if err != nil {
			return false
		}

		// Assume we have StartTransientUnit
		hasStartTransientUnit = true

		// But if we get UnknownMethod error we don't
		if _, err := theConn.StartTransientUnit("test.scope", "invalid"); err != nil {
			if dbusError, ok := err.(dbus.Error); ok {
				if dbusError.Name == "org.freedesktop.DBus.Error.UnknownMethod" {
					hasStartTransientUnit = false
					return hasStartTransientUnit
				}
			}
		}

		// Ensure the scope name we use doesn't exist. Use the Pid to
		// avoid collisions between multiple libcontainer users on a
		// single host.
		scope := fmt.Sprintf("libcontainer-%d-systemd-test-default-dependencies.scope", os.Getpid())
		testScopeExists := true
		for i := 0; i <= testScopeWait; i++ {
			if _, err := theConn.StopUnit(scope, "replace"); err != nil {
				if dbusError, ok := err.(dbus.Error); ok {
					if strings.Contains(dbusError.Name, "org.freedesktop.systemd1.NoSuchUnit") {
						testScopeExists = false
						break
					}
				}
			}
			time.Sleep(time.Millisecond)
		}

		// Bail out if we can't kill this scope without testing for DefaultDependencies
		if testScopeExists {
			return hasStartTransientUnit
		}

		// Assume StartTransientUnit on a scope allows DefaultDependencies
		hasTransientDefaultDependencies = true
		ddf := newProp("DefaultDependencies", false)
		if _, err := theConn.StartTransientUnit(scope, "replace", ddf); err != nil {
			if dbusError, ok := err.(dbus.Error); ok {
				if strings.Contains(dbusError.Name, "org.freedesktop.DBus.Error.PropertyReadOnly") {
					hasTransientDefaultDependencies = false
				}
			}
		}

		// Not critical because of the stop unit logic above.
		theConn.StopUnit(scope, "replace")
	}
	return hasStartTransientUnit
}

func getIfaceForUnit(unitName string) string {
	if strings.HasSuffix(unitName, ".scope") {
		return "Scope"
	}
	if strings.HasSuffix(unitName, ".service") {
		return "Service"
	}
	return "Unit"
}

func (m *Manager) Apply(pid int) error {
	var (
		c          = m.Cgroups
		unitName   = getUnitName(c)
		slice      = "system.slice"
		properties []systemd.Property
	)

	if c.Slice != "" {
		slice = c.Slice
	}

	properties = append(properties,
		systemd.PropSlice(slice),
		systemd.PropDescription("docker container "+c.Name),
		newProp("PIDs", []uint32{uint32(pid)}),
	)

	// Always enable accounting, this gets us the same behaviour as the fs implementation,
	// plus the kernel has some problems with joining the memory cgroup at a later time.
	properties = append(properties,
		newProp("MemoryAccounting", true),
		newProp("CPUAccounting", true),
		newProp("BlockIOAccounting", true))

	if hasTransientDefaultDependencies {
		properties = append(properties,
			newProp("DefaultDependencies", false))
	}

	if c.Memory != 0 {
		properties = append(properties,
			newProp("MemoryLimit", uint64(c.Memory)))
	}
	// TODO: MemoryReservation and MemorySwap not available in systemd

	if c.CpuShares != 0 {
		properties = append(properties,
			newProp("CPUShares", uint64(c.CpuShares)))
	}

	if c.BlkioWeight != 0 {
		properties = append(properties,
			newProp("BlockIOWeight", uint64(c.BlkioWeight)))
	}

	if _, err := theConn.StartTransientUnit(unitName, "replace", properties...); err != nil {
		return err
	}

	if err := joinDevices(c, pid); err != nil {
		return err
	}

	// TODO: CpuQuota and CpuPeriod not available in systemd
	// we need to manually join the cpu.cfs_quota_us and cpu.cfs_period_us
	if err := joinCpu(c, pid); err != nil {
		return err
	}

	// -1 disables memorySwap
	if c.MemorySwap >= 0 && c.Memory != 0 {
		if err := joinMemory(c, pid); err != nil {
			return err
		}

	}

	// we need to manually join the freezer and cpuset cgroup in systemd
	// because it does not currently support it via the dbus api.
	if err := joinFreezer(c, pid); err != nil {
		return err
	}

	if err := joinCpuset(c, pid); err != nil {
		return err
	}

	paths := make(map[string]string)
	for sysname := range subsystems {
		subsystemPath, err := getSubsystemPath(m.Cgroups, sysname)
		if err != nil {
			// Don't fail if a cgroup hierarchy was not found, just skip this subsystem
			if cgroups.IsNotFound(err) {
				continue
			}
			return err
		}
		paths[sysname] = subsystemPath
	}

	m.Paths = paths

	return nil
}

func (m *Manager) Destroy() error {
	return cgroups.RemovePaths(m.Paths)
}

func (m *Manager) GetPaths() map[string]string {
	return m.Paths
}

func writeFile(dir, file, data string) error {
	return ioutil.WriteFile(filepath.Join(dir, file), []byte(data), 0700)
}

func join(c *configs.Cgroup, subsystem string, pid int) (string, error) {
	path, err := getSubsystemPath(c, subsystem)
	if err != nil {
		return "", err
	}
	if err := os.MkdirAll(path, 0755); err != nil && !os.IsExist(err) {
		return "", err
	}
	if err := writeFile(path, "cgroup.procs", strconv.Itoa(pid)); err != nil {
		return "", err
	}

	return path, nil
}

func joinCpu(c *configs.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "cpu")
	if err != nil {
		return err
	}
	if c.CpuQuota != 0 {
		if err = ioutil.WriteFile(filepath.Join(path, "cpu.cfs_quota_us"), []byte(strconv.FormatInt(c.CpuQuota, 10)), 0700); err != nil {
			return err
		}
	}
	if c.CpuPeriod != 0 {
		if err = ioutil.WriteFile(filepath.Join(path, "cpu.cfs_period_us"), []byte(strconv.FormatInt(c.CpuPeriod, 10)), 0700); err != nil {
			return err
		}
	}
	return nil
}

func joinFreezer(c *configs.Cgroup, pid int) error {
	if _, err := join(c, "freezer", pid); err != nil {
		return err
	}

	return nil
}

func getSubsystemPath(c *configs.Cgroup, subsystem string) (string, error) {
	mountpoint, err := cgroups.FindCgroupMountpoint(subsystem)
	if err != nil {
		return "", err
	}

	initPath, err := cgroups.GetInitCgroupDir(subsystem)
	if err != nil {
		return "", err
	}

	slice := "system.slice"
	if c.Slice != "" {
		slice = c.Slice
	}

	return filepath.Join(mountpoint, initPath, slice, getUnitName(c)), nil
}

func (m *Manager) Freeze(state configs.FreezerState) error {
	path, err := getSubsystemPath(m.Cgroups, "freezer")
	if err != nil {
		return err
	}

	prevState := m.Cgroups.Freezer
	m.Cgroups.Freezer = state

	freezer := subsystems["freezer"]
	err = freezer.Set(path, m.Cgroups)
	if err != nil {
		m.Cgroups.Freezer = prevState
		return err
	}

	return nil
}

func (m *Manager) GetPids() ([]int, error) {
	path, err := getSubsystemPath(m.Cgroups, "cpu")
	if err != nil {
		return nil, err
	}

	return cgroups.ReadProcsFile(path)
}

func (m *Manager) GetStats() (*cgroups.Stats, error) {
	stats := cgroups.NewStats()
	for name, path := range m.Paths {
		sys, ok := subsystems[name]
		if !ok || !cgroups.PathExists(path) {
			continue
		}
		if err := sys.GetStats(path, stats); err != nil {
			return nil, err
		}
	}

	return stats, nil
}

func (m *Manager) Set(container *configs.Config) error {
	panic("not implemented")
}

func getUnitName(c *configs.Cgroup) string {
	return fmt.Sprintf("%s-%s.scope", c.Parent, c.Name)
}

// Atm we can't use the systemd device support because of two missing things:
// * Support for wildcards to allow mknod on any device
// * Support for wildcards to allow /dev/pts support
//
// The second is available in more recent systemd as "char-pts", but not in e.g. v208 which is
// in wide use. When both these are availalable we will be able to switch, but need to keep the old
// implementation for backwards compat.
//
// Note: we can't use systemd to set up the initial limits, and then change the cgroup
// because systemd will re-write the device settings if it needs to re-apply the cgroup context.
// This happens at least for v208 when any sibling unit is started.
func joinDevices(c *configs.Cgroup, pid int) error {
	path, err := join(c, "devices", pid)
	if err != nil {
		return err
	}

	devices := subsystems["devices"]
	if err := devices.Set(path, c); err != nil {
		return err
	}

	return nil
}

// Symmetrical public function to update device based cgroups.  Also available
// in the fs implementation.
func ApplyDevices(c *configs.Cgroup, pid int) error {
	return joinDevices(c, pid)
}

func joinMemory(c *configs.Cgroup, pid int) error {
	memorySwap := c.MemorySwap

	if memorySwap == 0 {
		// By default, MemorySwap is set to twice the size of RAM.
		memorySwap = c.Memory * 2
	}

	path, err := getSubsystemPath(c, "memory")
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filepath.Join(path, "memory.memsw.limit_in_bytes"), []byte(strconv.FormatInt(memorySwap, 10)), 0700)
}

// systemd does not atm set up the cpuset controller, so we must manually
// join it. Additionally that is a very finicky controller where each
// level must have a full setup as the default for a new directory is "no cpus"
func joinCpuset(c *configs.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "cpuset")
	if err != nil {
		return err
	}

	s := &fs.CpusetGroup{}

	return s.ApplyDir(path, c, pid)
}
<commit_msg>--crid=185625 - Disable using systemd to control cgroup settings<commit_after>// +build linux

package systemd

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	systemd "github.com/coreos/go-systemd/dbus"
	"github.com/docker/libcontainer/cgroups"
	"github.com/docker/libcontainer/cgroups/fs"
	"github.com/docker/libcontainer/configs"
	"github.com/godbus/dbus"
)

type Manager struct {
	Cgroups *configs.Cgroup
	Paths   map[string]string
}

type subsystem interface {
	// Returns the stats, as 'stats', corresponding to the cgroup under 'path'.
	GetStats(path string, stats *cgroups.Stats) error
	// Set the cgroup represented by cgroup.
	Set(path string, cgroup *configs.Cgroup) error
}

var subsystems = map[string]subsystem{
	"devices":    &fs.DevicesGroup{},
	"memory":     &fs.MemoryGroup{},
	"cpu":        &fs.CpuGroup{},
	"cpuset":     &fs.CpusetGroup{},
	"cpuacct":    &fs.CpuacctGroup{},
	"blkio":      &fs.BlkioGroup{},
	"perf_event": &fs.PerfEventGroup{},
	"freezer":    &fs.FreezerGroup{},
}

const (
	testScopeWait = 4
)

var (
	connLock                        sync.Mutex
	theConn                         *systemd.Conn
	hasStartTransientUnit           bool
	hasTransientDefaultDependencies bool
)

func newProp(name string, units interface{}) systemd.Property {
	return systemd.Property{
		Name:  name,
		Value: dbus.MakeVariant(units),
	}
}

func UseSystemd() bool {
	s, err := os.Stat("/run/systemd/system")
	if err != nil || !s.IsDir() {
		return false
	}

	connLock.Lock()
	defer connLock.Unlock()

	if theConn == nil {
		var err error
		theConn, err = systemd.New()
		if err != nil {
			return false
		}

		// Assume we have StartTransientUnit
		hasStartTransientUnit = true

		// But if we get UnknownMethod error we don't
		if _, err := theConn.StartTransientUnit("test.scope", "invalid"); err != nil {
			if dbusError, ok := err.(dbus.Error); ok {
				if dbusError.Name == "org.freedesktop.DBus.Error.UnknownMethod" {
					hasStartTransientUnit = false
					return hasStartTransientUnit
				}
			}
		}

		// Ensure the scope name we use doesn't exist. Use the Pid to
		// avoid collisions between multiple libcontainer users on a
		// single host.
		scope := fmt.Sprintf("libcontainer-%d-systemd-test-default-dependencies.scope", os.Getpid())
		testScopeExists := true
		for i := 0; i <= testScopeWait; i++ {
			if _, err := theConn.StopUnit(scope, "replace"); err != nil {
				if dbusError, ok := err.(dbus.Error); ok {
					if strings.Contains(dbusError.Name, "org.freedesktop.systemd1.NoSuchUnit") {
						testScopeExists = false
						break
					}
				}
			}
			time.Sleep(time.Millisecond)
		}

		// Bail out if we can't kill this scope without testing for DefaultDependencies
		if testScopeExists {
			return hasStartTransientUnit
		}

		// Assume StartTransientUnit on a scope allows DefaultDependencies
		hasTransientDefaultDependencies = true
		ddf := newProp("DefaultDependencies", false)
		if _, err := theConn.StartTransientUnit(scope, "replace", ddf); err != nil {
			if dbusError, ok := err.(dbus.Error); ok {
				if strings.Contains(dbusError.Name, "org.freedesktop.DBus.Error.PropertyReadOnly") {
					hasTransientDefaultDependencies = false
				}
			}
		}

		// Not critical because of the stop unit logic above.
		theConn.StopUnit(scope, "replace")
	}
//	return hasStartTransientUnit
	return false
}

func getIfaceForUnit(unitName string) string {
	if strings.HasSuffix(unitName, ".scope") {
		return "Scope"
	}
	if strings.HasSuffix(unitName, ".service") {
		return "Service"
	}
	return "Unit"
}

func (m *Manager) Apply(pid int) error {
	var (
		c          = m.Cgroups
		unitName   = getUnitName(c)
		slice      = "system.slice"
		properties []systemd.Property
	)

	if c.Slice != "" {
		slice = c.Slice
	}

	properties = append(properties,
		systemd.PropSlice(slice),
		systemd.PropDescription("docker container "+c.Name),
		newProp("PIDs", []uint32{uint32(pid)}),
	)

	// Always enable accounting, this gets us the same behaviour as the fs implementation,
	// plus the kernel has some problems with joining the memory cgroup at a later time.
	properties = append(properties,
		newProp("MemoryAccounting", true),
		newProp("CPUAccounting", true),
		newProp("BlockIOAccounting", true))

	if hasTransientDefaultDependencies {
		properties = append(properties,
			newProp("DefaultDependencies", false))
	}

	if c.Memory != 0 {
		properties = append(properties,
			newProp("MemoryLimit", uint64(c.Memory)))
	}
	// TODO: MemoryReservation and MemorySwap not available in systemd

	if c.CpuShares != 0 {
		properties = append(properties,
			newProp("CPUShares", uint64(c.CpuShares)))
	}

	if c.BlkioWeight != 0 {
		properties = append(properties,
			newProp("BlockIOWeight", uint64(c.BlkioWeight)))
	}

	if _, err := theConn.StartTransientUnit(unitName, "replace", properties...); err != nil {
		return err
	}

	if err := joinDevices(c, pid); err != nil {
		return err
	}

	// TODO: CpuQuota and CpuPeriod not available in systemd
	// we need to manually join the cpu.cfs_quota_us and cpu.cfs_period_us
	if err := joinCpu(c, pid); err != nil {
		return err
	}

	// -1 disables memorySwap
	if c.MemorySwap >= 0 && c.Memory != 0 {
		if err := joinMemory(c, pid); err != nil {
			return err
		}

	}

	// we need to manually join the freezer and cpuset cgroup in systemd
	// because it does not currently support it via the dbus api.
	if err := joinFreezer(c, pid); err != nil {
		return err
	}

	if err := joinCpuset(c, pid); err != nil {
		return err
	}

	paths := make(map[string]string)
	for sysname := range subsystems {
		subsystemPath, err := getSubsystemPath(m.Cgroups, sysname)
		if err != nil {
			// Don't fail if a cgroup hierarchy was not found, just skip this subsystem
			if cgroups.IsNotFound(err) {
				continue
			}
			return err
		}
		paths[sysname] = subsystemPath
	}

	m.Paths = paths

	return nil
}

func (m *Manager) Destroy() error {
	return cgroups.RemovePaths(m.Paths)
}

func (m *Manager) GetPaths() map[string]string {
	return m.Paths
}

func writeFile(dir, file, data string) error {
	return ioutil.WriteFile(filepath.Join(dir, file), []byte(data), 0700)
}

func join(c *configs.Cgroup, subsystem string, pid int) (string, error) {
	path, err := getSubsystemPath(c, subsystem)
	if err != nil {
		return "", err
	}
	if err := os.MkdirAll(path, 0755); err != nil && !os.IsExist(err) {
		return "", err
	}
	if err := writeFile(path, "cgroup.procs", strconv.Itoa(pid)); err != nil {
		return "", err
	}

	return path, nil
}

func joinCpu(c *configs.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "cpu")
	if err != nil {
		return err
	}
	if c.CpuQuota != 0 {
		if err = ioutil.WriteFile(filepath.Join(path, "cpu.cfs_quota_us"), []byte(strconv.FormatInt(c.CpuQuota, 10)), 0700); err != nil {
			return err
		}
	}
	if c.CpuPeriod != 0 {
		if err = ioutil.WriteFile(filepath.Join(path, "cpu.cfs_period_us"), []byte(strconv.FormatInt(c.CpuPeriod, 10)), 0700); err != nil {
			return err
		}
	}
	return nil
}

func joinFreezer(c *configs.Cgroup, pid int) error {
	if _, err := join(c, "freezer", pid); err != nil {
		return err
	}

	return nil
}

func getSubsystemPath(c *configs.Cgroup, subsystem string) (string, error) {
	mountpoint, err := cgroups.FindCgroupMountpoint(subsystem)
	if err != nil {
		return "", err
	}

	initPath, err := cgroups.GetInitCgroupDir(subsystem)
	if err != nil {
		return "", err
	}

	slice := "system.slice"
	if c.Slice != "" {
		slice = c.Slice
	}

	return filepath.Join(mountpoint, initPath, slice, getUnitName(c)), nil
}

func (m *Manager) Freeze(state configs.FreezerState) error {
	path, err := getSubsystemPath(m.Cgroups, "freezer")
	if err != nil {
		return err
	}

	prevState := m.Cgroups.Freezer
	m.Cgroups.Freezer = state

	freezer := subsystems["freezer"]
	err = freezer.Set(path, m.Cgroups)
	if err != nil {
		m.Cgroups.Freezer = prevState
		return err
	}

	return nil
}

func (m *Manager) GetPids() ([]int, error) {
	path, err := getSubsystemPath(m.Cgroups, "cpu")
	if err != nil {
		return nil, err
	}

	return cgroups.ReadProcsFile(path)
}

func (m *Manager) GetStats() (*cgroups.Stats, error) {
	stats := cgroups.NewStats()
	for name, path := range m.Paths {
		sys, ok := subsystems[name]
		if !ok || !cgroups.PathExists(path) {
			continue
		}
		if err := sys.GetStats(path, stats); err != nil {
			return nil, err
		}
	}

	return stats, nil
}

func (m *Manager) Set(container *configs.Config) error {
	panic("not implemented")
}

func getUnitName(c *configs.Cgroup) string {
	return fmt.Sprintf("%s-%s.scope", c.Parent, c.Name)
}

// Atm we can't use the systemd device support because of two missing things:
// * Support for wildcards to allow mknod on any device
// * Support for wildcards to allow /dev/pts support
//
// The second is available in more recent systemd as "char-pts", but not in e.g. v208 which is
// in wide use. When both these are availalable we will be able to switch, but need to keep the old
// implementation for backwards compat.
//
// Note: we can't use systemd to set up the initial limits, and then change the cgroup
// because systemd will re-write the device settings if it needs to re-apply the cgroup context.
// This happens at least for v208 when any sibling unit is started.
func joinDevices(c *configs.Cgroup, pid int) error {
	path, err := join(c, "devices", pid)
	if err != nil {
		return err
	}

	devices := subsystems["devices"]
	if err := devices.Set(path, c); err != nil {
		return err
	}

	return nil
}

// Symmetrical public function to update device based cgroups.  Also available
// in the fs implementation.
func ApplyDevices(c *configs.Cgroup, pid int) error {
	return joinDevices(c, pid)
}

func joinMemory(c *configs.Cgroup, pid int) error {
	memorySwap := c.MemorySwap

	if memorySwap == 0 {
		// By default, MemorySwap is set to twice the size of RAM.
		memorySwap = c.Memory * 2
	}

	path, err := getSubsystemPath(c, "memory")
	if err != nil {
		return err
	}

	return ioutil.WriteFile(filepath.Join(path, "memory.memsw.limit_in_bytes"), []byte(strconv.FormatInt(memorySwap, 10)), 0700)
}

// systemd does not atm set up the cpuset controller, so we must manually
// join it. Additionally that is a very finicky controller where each
// level must have a full setup as the default for a new directory is "no cpus"
func joinCpuset(c *configs.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "cpuset")
	if err != nil {
		return err
	}

	s := &fs.CpusetGroup{}

	return s.ApplyDir(path, c, pid)
}
<|endoftext|>


<commit_before>package service

import (
	"errors"
	"github.com/timeredbull/tsuru/api/app"
	"github.com/timeredbull/tsuru/db"
	"labix.org/v2/mgo/bson"
)

type App app.App

type ServiceInstance struct {
	Name        string `bson:"_id"`
	ServiceName string `bson:"service_name"`
	Apps        []string
	Teams       []string
	Instance    string
	Host        string
	State       string
	Env         map[string]string
}

func (si *ServiceInstance) Create() error {
	if si.State == "" {
		si.State = "creating"
	}
	return db.Session.ServiceInstances().Insert(si)
}

func (si *ServiceInstance) Delete() error {
	doc := bson.M{"_id": si.Name, "apps": si.Apps}
	return db.Session.ServiceInstances().Remove(doc)
}

func (si *ServiceInstance) Service() *Service {
	s := &Service{}
	db.Session.Services().Find(bson.M{"_id": si.ServiceName}).One(s)
	return s
}

func (si *ServiceInstance) AllApps() []app.App {
	var apps []app.App
	q := bson.M{"name": bson.M{"$in": si.Apps}}
	db.Session.Apps().Find(q).All(&apps)
	return apps
}

func (si *ServiceInstance) RemoveApp(appName string) error {
	index := -1
	for i, name := range si.Apps {
		if name == appName {
			index = i
			break
		}
	}
	if index < 0 {
		return errors.New("This app is not binded to this service instance.")
	}
	last := len(si.Apps) - 1
	if index != last {
		si.Apps[index] = si.Apps[last]
	}
	si.Apps = si.Apps[:last]
	return nil
}
<commit_msg>Revert "api/service: added type App to the service"<commit_after>package service

import (
	"errors"
	"github.com/timeredbull/tsuru/api/app"
	"github.com/timeredbull/tsuru/db"
	"labix.org/v2/mgo/bson"
)

type ServiceInstance struct {
	Name        string `bson:"_id"`
	ServiceName string `bson:"service_name"`
	Apps        []string
	Teams       []string
	Instance    string
	Host        string
	State       string
	Env         map[string]string
}

func (si *ServiceInstance) Create() error {
	if si.State == "" {
		si.State = "creating"
	}
	return db.Session.ServiceInstances().Insert(si)
}

func (si *ServiceInstance) Delete() error {
	doc := bson.M{"_id": si.Name, "apps": si.Apps}
	return db.Session.ServiceInstances().Remove(doc)
}

func (si *ServiceInstance) Service() *Service {
	s := &Service{}
	db.Session.Services().Find(bson.M{"_id": si.ServiceName}).One(s)
	return s
}

func (si *ServiceInstance) AllApps() []app.App {
	var apps []app.App
	q := bson.M{"name": bson.M{"$in": si.Apps}}
	db.Session.Apps().Find(q).All(&apps)
	return apps
}

func (si *ServiceInstance) RemoveApp(appName string) error {
	index := -1
	for i, name := range si.Apps {
		if name == appName {
			index = i
			break
		}
	}
	if index < 0 {
		return errors.New("This app is not binded to this service instance.")
	}
	last := len(si.Apps) - 1
	if index != last {
		si.Apps[index] = si.Apps[last]
	}
	si.Apps = si.Apps[:last]
	return nil
}
<|endoftext|>


<commit_before>package util

import (
	"fmt"
	"os/exec"
	"strings"
	"time"

	"k8s.io/kubernetes/pkg/util/wait"
)

// Database interface allows testing database images.
type Database interface {
	// PodName returns the name of the Pod this helper is bound to.
	PodName() string

	// IsReady indicates whether the underlying Pod is ready for queries.
	IsReady(oc *CLI) (bool, error)

	// Query queries the database as a regular user.
	Query(oc *CLI, query string) (string, error)

	// QueryPrivileged queries the database as a privileged user.
	QueryPrivileged(oc *CLI, query string) (string, error)

	// TestRemoteLogin tests whether it is possible to remote login to hostAddress.
	TestRemoteLogin(oc *CLI, hostAddress string) error
}

// ReplicaSet interface allows to interact with database on multiple nodes.
type ReplicaSet interface {
	// QueryPrimary queries the database on primary node as a regular user.
	QueryPrimary(oc *CLI, query string) (string, error)
}

// WaitForQueryOutputSatisfies will execute the query multiple times, until the
// specified predicate function is return true.
func WaitForQueryOutputSatisfies(oc *CLI, d Database, timeout time.Duration, admin bool, query string, predicate func(string) bool) error {
	err := wait.Poll(5*time.Second, timeout, func() (bool, error) {
		var (
			out string
			err error
		)

		if admin {
			out, err = d.QueryPrivileged(oc, query)
		} else {
			out, err = d.Query(oc, query)
		}
		if _, ok := err.(*exec.ExitError); ok {
			// Ignore exit errors
			return false, nil
		}
		if err != nil {
			return false, err
		}
		if predicate(out) {
			return true, nil
		}
		return false, nil
	})
	if err == wait.ErrWaitTimeout {
		return fmt.Errorf("timed out waiting for query: %q", query)
	}
	return err
}

// WaitForQueryOutputContains will execute the query multiple times, until the
// specified substring is found in the results. This function should be used for
// testing replication, since it might take some time until the data is propagated
// to slaves.
func WaitForQueryOutputContains(oc *CLI, d Database, timeout time.Duration, admin bool, query, resultSubstr string) error {
	return WaitForQueryOutputSatisfies(oc, d, timeout, admin, query, func(resultOutput string) bool {
		return strings.Contains(resultOutput, resultSubstr)
	})
}

// WaitUntilUp continuously waits for the server to become ready, up until timeout.
func WaitUntilUp(oc *CLI, d Database, timeout time.Duration) error {
	err := wait.Poll(2*time.Second, timeout, func() (bool, error) {
		return d.IsReady(oc)
	})
	if err == wait.ErrWaitTimeout {
		return fmt.Errorf("timed out waiting for pod %s get up", d.PodName())
	}
	return err
}

// WaitUntilAllHelpersAreUp waits until all helpers are ready to serve requests.
func WaitUntilAllHelpersAreUp(oc *CLI, helpers []Database) error {
	for _, m := range helpers {
		if err := WaitUntilUp(oc, m, 3*time.Minute); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>check for correct exiterr type on oc exec failures<commit_after>package util

import (
	"fmt"
	"reflect"
	"strings"
	"time"

	"k8s.io/kubernetes/pkg/util/wait"
	e2e "k8s.io/kubernetes/test/e2e/framework"
)

// Database interface allows testing database images.
type Database interface {
	// PodName returns the name of the Pod this helper is bound to.
	PodName() string

	// IsReady indicates whether the underlying Pod is ready for queries.
	IsReady(oc *CLI) (bool, error)

	// Query queries the database as a regular user.
	Query(oc *CLI, query string) (string, error)

	// QueryPrivileged queries the database as a privileged user.
	QueryPrivileged(oc *CLI, query string) (string, error)

	// TestRemoteLogin tests whether it is possible to remote login to hostAddress.
	TestRemoteLogin(oc *CLI, hostAddress string) error
}

// ReplicaSet interface allows to interact with database on multiple nodes.
type ReplicaSet interface {
	// QueryPrimary queries the database on primary node as a regular user.
	QueryPrimary(oc *CLI, query string) (string, error)
}

// WaitForQueryOutputSatisfies will execute the query multiple times, until the
// specified predicate function is return true.
func WaitForQueryOutputSatisfies(oc *CLI, d Database, timeout time.Duration, admin bool, query string, predicate func(string) bool) error {
	err := wait.Poll(5*time.Second, timeout, func() (bool, error) {
		var (
			out string
			err error
		)

		if admin {
			out, err = d.QueryPrivileged(oc, query)
		} else {
			out, err = d.Query(oc, query)
		}
		if _, ok := err.(*ExitError); ok {
			// Ignore exit errors
			return false, nil
		}
		if err != nil {
			e2e.Logf("failing immediately with error: %v, type=%v", err, reflect.TypeOf(err))
			return false, err
		}
		if predicate(out) {
			return true, nil
		}
		return false, nil
	})
	if err == wait.ErrWaitTimeout {
		return fmt.Errorf("timed out waiting for query: %q", query)
	}
	return err
}

// WaitForQueryOutputContains will execute the query multiple times, until the
// specified substring is found in the results. This function should be used for
// testing replication, since it might take some time until the data is propagated
// to slaves.
func WaitForQueryOutputContains(oc *CLI, d Database, timeout time.Duration, admin bool, query, resultSubstr string) error {
	return WaitForQueryOutputSatisfies(oc, d, timeout, admin, query, func(resultOutput string) bool {
		return strings.Contains(resultOutput, resultSubstr)
	})
}

// WaitUntilUp continuously waits for the server to become ready, up until timeout.
func WaitUntilUp(oc *CLI, d Database, timeout time.Duration) error {
	err := wait.Poll(2*time.Second, timeout, func() (bool, error) {
		return d.IsReady(oc)
	})
	if err == wait.ErrWaitTimeout {
		return fmt.Errorf("timed out waiting for pod %s get up", d.PodName())
	}
	return err
}

// WaitUntilAllHelpersAreUp waits until all helpers are ready to serve requests.
func WaitUntilAllHelpersAreUp(oc *CLI, helpers []Database) error {
	for _, m := range helpers {
		if err := WaitUntilUp(oc, m, 3*time.Minute); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>