<commit_before>package movavg_test

import (
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

func TestMovavg(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Movavg Suite")
}

// ----------------------------------------------------------------------------

type mockMA struct {
	addArgs []float64
	addRes  float64
	avgRes  float64
}

func (a *mockMA) Add(v float64) float64 {
	a.addArgs = append(a.addArgs, v)
	return a.addRes
}

func (a *mockMA) Avg() float64 {
	return a.avgRes
}

// ----------------------------------------------------------------------------

type mockSet struct {
	addArgs []float64
	addRes  []float64
	avgRes  []float64
}

func (s *mockSet) Add(v float64) []float64 {
	s.addArgs = append(s.addArgs, v)
	return s.addRes
}

func (s *mockSet) Avg() []float64 {
	return s.avgRes
}
<commit_msg>multi: bring some type assertions for mocks<commit_after>package movavg_test

import (
	"testing"

	. "github.com/mxmCherry/movavg"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

func TestMovavg(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Movavg Suite")
}

// ----------------------------------------------------------------------------

var _ MA = (*mockMA)(nil)

type mockMA struct {
	addArgs []float64
	addRes  float64
	avgRes  float64
}

func (a *mockMA) Add(v float64) float64 {
	a.addArgs = append(a.addArgs, v)
	return a.addRes
}

func (a *mockMA) Avg() float64 {
	return a.avgRes
}

// ----------------------------------------------------------------------------

var _ MultiMA = (*mockSet)(nil)

type mockSet struct {
	addArgs []float64
	addRes  []float64
	avgRes  []float64
}

func (s *mockSet) Add(v float64) []float64 {
	s.addArgs = append(s.addArgs, v)
	return s.addRes
}

func (s *mockSet) Avg() []float64 {
	return s.avgRes
}
<|endoftext|>


<commit_before>package goutils

import (
	"encoding/json"
	"encoding/xml"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"

	"github.com/hoveychen/go-utils/flags"
	"github.com/pkg/errors"

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

var (
	proxyAddr      = flags.String("proxy", "", "Specify proxy address to fetch data")
	proxyType      = flags.String("proxyType", "sock5", "Either sock5 or http for proxy.")
	requestTimeout = flags.Int("requestTimeout", 10, "Timeout in sec when fetching a remote page.")
	downloadClient *http.Client
)

func init() {
	PkgInit(func() {
		httpTransport := &http.Transport{}
		if *proxyAddr != "" {
			switch *proxyType {
			case "http":
				proxyUrl, err := url.Parse(*proxyAddr)
				if err != nil {
					LogFatal("Failed to parse --proxyAddr", err)
				}
				httpTransport.Proxy = http.ProxyURL(proxyUrl)
			case "sock5":
				dialer, err := proxy.SOCKS5("tcp", *proxyAddr, nil, proxy.Direct)
				if err != nil {
					LogFatal("Failed to dial sock5", err)
				}
				httpTransport.Dial = dialer.Dial
			default:
				LogFatal("Unknown proxy type:", *proxyType)
			}
		}

		downloadClient = &http.Client{Transport: httpTransport}
		downloadClient.Timeout = time.Duration(*requestTimeout) * time.Second
	})
}

// FetchData is a helper function to load local/remote data in the same function.
// Local: goutils.FetchData("/absolute/path/to/file")
// Remote: goutils.FetchData("https://www.google.com")
// Also, it's integrated with proxy in flags.
// TODO(yuheng): Allow more options, while keeping easy use.
func FetchData(path string) ([]byte, error) {
	url, err := url.Parse(path)
	if err != nil {
		return nil, errors.Wrap(err, "decode url")
	}

	switch url.Scheme {
	case "http", "https":
		req, err := http.NewRequest("GET", path, nil)
		if err != nil {
			return nil, errors.Wrap(err, "new request")
		}
		resp, err := downloadClient.Do(req)
		if err != nil {
			return nil, errors.Wrap(err, "do request")
		}
		defer resp.Body.Close()
		data, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, errors.Wrap(err, "read response")
		}

		return data, nil
	case "":
		data, err := ioutil.ReadFile(url.Path)
		if err != nil {
			return nil, errors.Wrap(err, "read file")
		}
		return data, nil
	default:
		return nil, errors.New("Unknown scheme:" + url.Scheme)
	}
}

// FetchJson is a wrapper to call FetchData() and parse results from json.
func FetchJson(path string, resp interface{}) error {
	d, err := FetchData(path)
	if err != nil {
		return errors.Wrap(err, "fetch data")
	}
	if err := json.Unmarshal(d, resp); err != nil {
		return errors.Wrap(err, "decode json")
	}
	return nil
}

// FetchXml is a wrapper to call FetchData() and parse results from xml.
func FetchXml(path string, resp interface{}) error {
	d, err := FetchData(path)
	if err != nil {
		return errors.Wrap(err, "fetch data")
	}
	if err := xml.Unmarshal(d, resp); err != nil {
		return errors.Wrap(err, "decode xml")
	}
	return nil
}
<commit_msg>[Request] Add redirect times.<commit_after>package goutils

import (
	"encoding/json"
	"encoding/xml"
	"io/ioutil"
	"net/http"
	"net/url"
	"time"

	"github.com/hoveychen/go-utils/flags"
	"github.com/pkg/errors"

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

var (
	proxyAddr      = flags.String("proxy", "", "Specify proxy address to fetch data")
	proxyType      = flags.String("proxyType", "sock5", "Either sock5 or http for proxy.")
	requestTimeout = flags.Int("requestTimeout", 10, "Timeout in sec when fetching a remote page.")
	downloadClient *http.Client
)

func init() {
	PkgInit(func() {
		httpTransport := &http.Transport{}
		if *proxyAddr != "" {
			switch *proxyType {
			case "http":
				proxyUrl, err := url.Parse(*proxyAddr)
				if err != nil {
					LogFatal("Failed to parse --proxyAddr", err)
				}
				httpTransport.Proxy = http.ProxyURL(proxyUrl)
			case "sock5":
				dialer, err := proxy.SOCKS5("tcp", *proxyAddr, nil, proxy.Direct)
				if err != nil {
					LogFatal("Failed to dial sock5", err)
				}
				httpTransport.Dial = dialer.Dial
			default:
				LogFatal("Unknown proxy type:", *proxyType)
			}
		}

		downloadClient = &http.Client{Transport: httpTransport}
		downloadClient.Timeout = time.Duration(*requestTimeout) * time.Second
		downloadClient.CheckRedirect = modifiedCheckRedirect
	})
}

func modifiedCheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= 15 {
		return errors.New("stopped after 15 redirects")
	}
	return nil
}

// FetchData is a helper function to load local/remote data in the same function.
// Local: goutils.FetchData("/absolute/path/to/file")
// Remote: goutils.FetchData("https://www.google.com")
// Also, it's integrated with proxy in flags.
// TODO(yuheng): Allow more options, while keeping easy use.
func FetchData(path string) ([]byte, error) {
	url, err := url.Parse(path)
	if err != nil {
		return nil, errors.Wrap(err, "decode url")
	}

	switch url.Scheme {
	case "http", "https":
		req, err := http.NewRequest("GET", path, nil)
		if err != nil {
			return nil, errors.Wrap(err, "new request")
		}
		resp, err := downloadClient.Do(req)
		if err != nil {
			return nil, errors.Wrap(err, "do request")
		}
		defer resp.Body.Close()
		data, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, errors.Wrap(err, "read response")
		}

		return data, nil
	case "":
		data, err := ioutil.ReadFile(url.Path)
		if err != nil {
			return nil, errors.Wrap(err, "read file")
		}
		return data, nil
	default:
		return nil, errors.New("Unknown scheme:" + url.Scheme)
	}
}

// FetchJson is a wrapper to call FetchData() and parse results from json.
func FetchJson(path string, resp interface{}) error {
	d, err := FetchData(path)
	if err != nil {
		return errors.Wrap(err, "fetch data")
	}
	if err := json.Unmarshal(d, resp); err != nil {
		return errors.Wrap(err, "decode json")
	}
	return nil
}

// FetchXml is a wrapper to call FetchData() and parse results from xml.
func FetchXml(path string, resp interface{}) error {
	d, err := FetchData(path)
	if err != nil {
		return errors.Wrap(err, "fetch data")
	}
	if err := xml.Unmarshal(d, resp); err != nil {
		return errors.Wrap(err, "decode xml")
	}
	return nil
}
<|endoftext|>


<commit_before>package movingmedian

import (
	"math"
	"math/rand"
	"testing"
)

func TestSameNumberInBothHeaps(t *testing.T) {
	tests := []struct {
		name       string
		windowSize int
		data       []float64
		want       []float64
	}{
		{
			"OneWindowSize",
			1,
			[]float64{1, 3, 5, 7, 9, 11, math.NaN()},
			[]float64{1, 3, 5, 7, 9, 11, math.NaN()},
		},
		{
			"OddWindowSize",
			3,
			[]float64{1, 3, 5, 7, 9, 11},
			[]float64{1, 2, 3, 5, 7, 9},
		},
		{
			"EvenWindowSize",
			4,
			[]float64{1, 3, 5, 7, 9, 11},
			[]float64{1, 2, 3, 4, 6, 8},
		},
		{
			"DecreasingValues",
			4,
			[]float64{19, 17, 15, 13, 11, 9},
			[]float64{19, 18, 17, 16, 14, 12},
		},
		{
			"SameNumberInBothHeaps3Times",
			4,
			[]float64{11, 13, 13, 13, 25, 27, 29, 31},
			[]float64{11, 12, 13, 13, 13, 19, 26, 28},
		},
		{
			"SameNumberInBothHeaps4Times",
			4,
			[]float64{11, 13, 13, 13, 13, 25, 27, 29, 31},
			[]float64{11, 12, 13, 13, 13, 13, 19, 26, 28},
		},
	}

	for _, test := range tests {
		m := NewMovingMedian(test.windowSize)
		for i, v := range test.data {
			m.Push(v)
			actual := m.Median()
			if test.want[i] != actual && !(math.IsNaN(actual) && math.IsNaN(test.want[i])) {
				firstElement := 1 + i - test.windowSize
				if firstElement < 0 {
					firstElement = 0
				}
				t.Errorf("failed on test %s index %d the median of %f is %f and not %f",
					test.name,
					i,
					test.data[firstElement:1+i],
					test.want[i],
					actual)
			}
		}
	}
}

func Benchmark_10values_windowsize1(b *testing.B) {
	benchmark(b, 10, 1)
}

func Benchmark_100values_windowsize10(b *testing.B) {
	benchmark(b, 100, 10)
}

func Benchmark_10Kvalues_windowsize100(b *testing.B) {
	benchmark(b, 10000, 100)
}

func Benchmark_10Kvalues_windowsize1000(b *testing.B) {
	benchmark(b, 10000, 1000)
}

func benchmark(b *testing.B, numberOfValues, windowSize int) {
	data := getData(numberOfValues)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		m := NewMovingMedian(windowSize)
		for _, v := range data {
			m.Push(v)
			m.Median()
		}
	}
}

func getData(rangeSize int) []float64 {
	var data = make([]float64, rangeSize)
	var r = rand.New(rand.NewSource(99))
	for i, _ := range data {
		data[i] = r.Float64()
	}

	return data
}
<commit_msg>added more unit tests, no errors found<commit_after>package movingmedian

import (
	"math"
	"math/rand"
	"testing"
)

func TestSameNumberInBothHeaps(t *testing.T) {
	tests := []struct {
		name       string
		windowSize int
		data       []float64
		want       []float64
	}{
		{
			"OneWindowSize",
			1,
			[]float64{1, 3, 5, 7, 9, 11, math.NaN()},
			[]float64{1, 3, 5, 7, 9, 11, math.NaN()},
		},
		{
			"OddWindowSize",
			3,
			[]float64{1, 3, 5, 7, 9, 11},
			[]float64{1, 2, 3, 5, 7, 9},
		},
		{
			"EvenWindowSize",
			4,
			[]float64{1, 3, 5, 7, 9, 11},
			[]float64{1, 2, 3, 4, 6, 8},
		},
		{
			"DecreasingValues",
			4,
			[]float64{19, 17, 15, 13, 11, 9},
			[]float64{19, 18, 17, 16, 14, 12},
		},
		{
			"DecreasingIncreasingValues",
			4,
			[]float64{21, 19, 17, 15, 13, 11, 13, 15, 17, 19},
			[]float64{21, 20, 19, 18, 16, 14, 13, 13, 14, 16},
		},
		{
			"IncreasingDecreasingValues",
			4,
			[]float64{11, 13, 15, 17, 19, 21, 19, 17, 15, 13},
			[]float64{11, 12, 13, 14, 16, 18, 19, 19, 18, 16},
		},
		{

			"ZigZag",
			4,
			[]float64{21, 23, 17, 27, 13, 31, 9, 35, 5, 39, 1},
			[]float64{21, 22, 21, 22, 20, 22, 20, 22, 20, 22, 20},
		},
		{

			"NewValuesInBetween",
			4,
			[]float64{21, 21, 19, 19, 21, 21, 19, 19, 19, 19},
			[]float64{21, 21, 21, 20, 20, 20, 20, 20, 19, 19},
		},
		{
			"SameNumberInBothHeaps3Times",
			4,
			[]float64{11, 13, 13, 13, 25, 27, 29, 31},
			[]float64{11, 12, 13, 13, 13, 19, 26, 28},
		},
		{
			"SameNumberInBothHeaps3TimesDecreasing",
			4,
			[]float64{31, 29, 29, 29, 17, 15, 13, 11},
			[]float64{31, 30, 29, 29, 29, 23, 16, 14},
		},
		{
			"SameNumberInBothHeaps4Times",
			4,
			[]float64{11, 13, 13, 13, 13, 25, 27, 29, 31},
			[]float64{11, 12, 13, 13, 13, 13, 19, 26, 28},
		},
	}

	for _, test := range tests {
		m := NewMovingMedian(test.windowSize)
		for i, v := range test.data {
			m.Push(v)
			actual := m.Median()
			if test.want[i] != actual && !(math.IsNaN(actual) && math.IsNaN(test.want[i])) {
				firstElement := 1 + i - test.windowSize
				if firstElement < 0 {
					firstElement = 0
				}
				t.Errorf("failed on test %s index %d the median of %f is %f and not %f",
					test.name,
					i,
					test.data[firstElement:1+i],
					test.want[i],
					actual)
			}
		}
	}
}

func Benchmark_10values_windowsize1(b *testing.B) {
	benchmark(b, 10, 1)
}

func Benchmark_100values_windowsize10(b *testing.B) {
	benchmark(b, 100, 10)
}

func Benchmark_10Kvalues_windowsize100(b *testing.B) {
	benchmark(b, 10000, 100)
}

func Benchmark_10Kvalues_windowsize1000(b *testing.B) {
	benchmark(b, 10000, 1000)
}

func benchmark(b *testing.B, numberOfValues, windowSize int) {
	data := getData(numberOfValues)

	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		m := NewMovingMedian(windowSize)
		for _, v := range data {
			m.Push(v)
			m.Median()
		}
	}
}

func getData(rangeSize int) []float64 {
	var data = make([]float64, rangeSize)
	var r = rand.New(rand.NewSource(99))
	for i, _ := range data {
		data[i] = r.Float64()
	}

	return data
}
<|endoftext|>


<commit_before>package maodou

import (
	"fmt"
	"log"
	"net/http"
	urlpkg "net/url"
	"time"

	"github.com/mnhkahn/maodou/request"
	"github.com/mnhkahn/maodou/request/goreq"
	"github.com/mnhkahn/maodou/request/proxy"
)

type Request struct {
	goreq.Request
	proxy    proxy.ProxyContainer
	root     string
	Interval time.Duration
}

func NewRequest(interval time.Duration) *Request {
	req := new(Request)
	req.Method = "GET"
	req.Timeout = time.Duration(30) * time.Second
	req.AddHeader("Accept-Language", "zh-CN,zh;q=0.8,en;q=0.6,zh-TW;q=0.4")
	req.AddHeader("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Interval = interval
	// req.ShowDebug = true

	goreq.SetConnectTimeout(5 * time.Second)

	return req
}

const (
	CAWL_NOPROXY = 0
	CAWL_PROXY   = 1
	CAWL_RETRY   = 2
)

func (this *Request) Cawl(paras ...interface{}) (*Response, error) {
	this.Uri = paras[0].(string)

	// Add referer
	if this.root == "" {
		this.root = this.Uri
	} else {
		this.UpdateHeader("Referer", this.root)
	}

	var p *proxy.ProxyConfig
	if len(paras) == 1 || (len(paras) == 2 && paras[1].(int) == CAWL_PROXY) {
		u := new(urlpkg.URL)
		p = this.proxy.One()
		if p.Ip != "" {
			u.Scheme = "http"
			u.Host = fmt.Sprintf("%s:%d", p.Ip, p.Port)
			this.Proxy = u.String()
		}
	}

	log.Println("Start to Parse: ", this.Uri)

	start := time.Now()
	this.ShowDebug = true
	this.UserAgent = request.UserAgent()
	http_resp, err := this.Do()
	log.Printf("Cawl use %v.\n", time.Now().Sub(start))
	// 修复代理错乱的问题，需要重置代理
	this.Proxy = ""
	if err != nil {
		if len(paras) == 1 || (len(paras) == 2 && paras[1].(int) == CAWL_PROXY) {
			this.proxy.DeleteProxy(p.Id)
		}
		log.Printf("Cawl Error: %s\n", err.Error())

		// Retry
		if len(paras) == 2 && paras[1].(int) == CAWL_RETRY {
			log.Println("Retry...")
			this.Cawl(paras...)
		} else {
			return nil, err
		}
	}

	var resp *Response
	if http_resp.StatusCode == http.StatusOK {
		resp, err = NewResponse(http_resp.Body, this.Uri)
		if err != nil {
			if len(paras) == 2 && paras[1].(int) == CAWL_RETRY {
				log.Println("Retry...")
				this.Cawl(paras...)
			} else {
				log.Printf("Cawl Error: %s.\n", err.Error())
				return resp, err
			}
		} else {
			log.Println("Cawl Success.")
		}
	} else {
		if len(paras) == 1 || (len(paras) == 2 && paras[1].(int) == CAWL_PROXY) {
			this.proxy.DeleteProxy(p.Id)
		}
		if len(paras) == 2 && paras[1].(int) == CAWL_RETRY {
			log.Println("Retry...")
			this.Cawl(paras...)
		} else {
			if http_resp.StatusCode == http.StatusMovedPermanently || http_resp.StatusCode == http.StatusFound {
				log.Println(this.Uri, http_resp.StatusCode)
				return this.Cawl(http_resp.Header.Get("Location"))
			} else {
				log.Printf("Cawl Got Status Code %d.\n", http_resp.StatusCode)
				return resp, fmt.Errorf("Cawl Got Status Code %d.", http_resp.StatusCode)
			}
		}
	}

	if this.Interval > 0 {
		time.Sleep(this.Interval)
	}
	return resp, nil
}

func (this *Request) InitProxy(proxy_name, proxy_config string) {
	this.proxy, _ = proxy.NewProxy(proxy_name, proxy_config)
	this.proxy.Init()
}

func (this *Request) DumpRequest() string {
	return this.Uri + "?" + urlpkg.Values(this.QueryString.(urlpkg.Values)).Encode()
}
<commit_msg>fix proxy nil bug<commit_after>package maodou

import (
	"fmt"
	"log"
	"net/http"
	urlpkg "net/url"
	"time"

	"github.com/mnhkahn/maodou/request"
	"github.com/mnhkahn/maodou/request/goreq"
	"github.com/mnhkahn/maodou/request/proxy"
)

type Request struct {
	goreq.Request
	proxy    proxy.ProxyContainer
	root     string
	Interval time.Duration
}

func NewRequest(interval time.Duration) *Request {
	req := new(Request)
	req.Method = "GET"
	req.Timeout = time.Duration(30) * time.Second
	req.AddHeader("Accept-Language", "zh-CN,zh;q=0.8,en;q=0.6,zh-TW;q=0.4")
	req.AddHeader("Accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8")
	req.Interval = interval
	// req.ShowDebug = true

	goreq.SetConnectTimeout(5 * time.Second)

	return req
}

const (
	CAWL_NOPROXY = 0
	CAWL_PROXY   = 1
	CAWL_RETRY   = 2
)

func (this *Request) Cawl(paras ...interface{}) (*Response, error) {
	this.Uri = paras[0].(string)

	// Add referer
	if this.root == "" {
		this.root = this.Uri
	} else {
		this.UpdateHeader("Referer", this.root)
	}

	var p *proxy.ProxyConfig
	if this.proxy != nil && len(paras) == 1 || (len(paras) == 2 && paras[1].(int) == CAWL_PROXY) {
		u := new(urlpkg.URL)
		p = this.proxy.One()
		if p.Ip != "" {
			u.Scheme = "http"
			u.Host = fmt.Sprintf("%s:%d", p.Ip, p.Port)
			this.Proxy = u.String()
		}
	}

	log.Println("Start to Parse: ", this.Uri)

	start := time.Now()
	this.ShowDebug = true
	this.UserAgent = request.UserAgent()
	http_resp, err := this.Do()
	log.Printf("Cawl use %v.\n", time.Now().Sub(start))
	// 修复代理错乱的问题，需要重置代理
	this.Proxy = ""
	if err != nil {
		if len(paras) == 1 || (len(paras) == 2 && paras[1].(int) == CAWL_PROXY) {
			this.proxy.DeleteProxy(p.Id)
		}
		log.Printf("Cawl Error: %s\n", err.Error())

		// Retry
		if len(paras) == 2 && paras[1].(int) == CAWL_RETRY {
			log.Println("Retry...")
			this.Cawl(paras...)
		} else {
			return nil, err
		}
	}

	var resp *Response
	if http_resp.StatusCode == http.StatusOK {
		resp, err = NewResponse(http_resp.Body, this.Uri)
		if err != nil {
			if len(paras) == 2 && paras[1].(int) == CAWL_RETRY {
				log.Println("Retry...")
				this.Cawl(paras...)
			} else {
				log.Printf("Cawl Error: %s.\n", err.Error())
				return resp, err
			}
		} else {
			log.Println("Cawl Success.")
		}
	} else {
		if len(paras) == 1 || (len(paras) == 2 && paras[1].(int) == CAWL_PROXY) {
			this.proxy.DeleteProxy(p.Id)
		}
		if len(paras) == 2 && paras[1].(int) == CAWL_RETRY {
			log.Println("Retry...")
			this.Cawl(paras...)
		} else {
			if http_resp.StatusCode == http.StatusMovedPermanently || http_resp.StatusCode == http.StatusFound {
				log.Println(this.Uri, http_resp.StatusCode)
				if len(paras) == 2 {
					return this.Cawl(http_resp.Header.Get("Location"), paras[1])
				}
				return this.Cawl(http_resp.Header.Get("Location"))
			} else {
				log.Printf("Cawl Got Status Code %d.\n", http_resp.StatusCode)
				return resp, fmt.Errorf("Cawl Got Status Code %d.", http_resp.StatusCode)
			}
		}
	}

	if this.Interval > 0 {
		time.Sleep(this.Interval)
	}
	return resp, nil
}

func (this *Request) InitProxy(proxy_name, proxy_config string) {
	this.proxy, _ = proxy.NewProxy(proxy_name, proxy_config)
	this.proxy.Init()
}

func (this *Request) DumpRequest() string {
	return this.Uri + "?" + urlpkg.Values(this.QueryString.(urlpkg.Values)).Encode()
}
<|endoftext|>


<commit_before>package request

import (
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"

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

const Version = "0.2.0"

type FileField struct {
	FieldName string
	FileName  string
	File      io.Reader
}

type BasicAuth struct {
	Username string
	Password string
}

type Args struct {
	Client    *http.Client
	Headers   map[string]string
	Cookies   map[string]string
	Data      map[string]string
	Params    map[string]string
	Files     []FileField
	Json      interface{}
	Proxy     string
	BasicAuth BasicAuth
}

type Request struct {
	*Args
}

func NewArgs(c *http.Client) *Args {
	if c.Jar == nil {
		options := cookiejar.Options{
			PublicSuffixList: publicsuffix.List,
		}
		jar, _ := cookiejar.New(&options)
		c.Jar = jar
	}

	return &Args{
		Client:    c,
		Headers:   defaultHeaders,
		Cookies:   nil,
		Data:      nil,
		Params:    nil,
		Files:     nil,
		Json:      nil,
		Proxy:     "",
		BasicAuth: BasicAuth{},
	}
}

func NewRequest(c *http.Client) *Request {
	return &Request{NewArgs(c)}
}

func newURL(u string, params map[string]string) string {
	if params == nil {
		return u
	}

	p := url.Values{}
	for k, v := range params {
		p.Set(k, v)
	}
	if strings.Contains(u, "?") {
		return u + "&" + p.Encode()
	}
	return u + "?" + p.Encode()
}

func newBody(a *Args) (body io.Reader, contentType string, err error) {
	if a.Data == nil && a.Files == nil && a.Json == nil {
		return nil, "", nil
	}
	if a.Files != nil {
		return newMultipartBody(a)
	} else if a.Json != nil {
		return newJsonBody(a)
	}

	d := url.Values{}
	for k, v := range a.Data {
		d.Set(k, v)
	}
	return strings.NewReader(d.Encode()), "", nil
}

func newRequest(method string, url string, a *Args) (resp *Response, err error) {
	body, contentType, err := newBody(a)
	u := newURL(url, a.Params)
	req, err := http.NewRequest(method, u, body)
	if err != nil {
		return nil, err
	}
	applyHeaders(a, req, contentType)
	applyCookies(a, req)
	applyProxy(a)
	applyCheckRdirect(a)

	if a.BasicAuth.Username != "" {
		req.SetBasicAuth(a.BasicAuth.Username, a.BasicAuth.Password)
	}

	s, err := a.Client.Do(req)
	resp = &Response{s, nil}
	return
}

// Get issues a GET to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Get(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("GET", url, a)
	return
}

func (req *Request) Get(url interface{}) (resp *Response, err error) {
	resp, err = Get(url2string(url), req2arg(req))
	return
}

// Head issues a HEAD to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Head(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("HEAD", url, a)
	return
}

func (req *Request) Head(url interface{}) (resp *Response, err error) {
	resp, err = Head(url2string(url), req2arg(req))
	return
}

// Post issues a POST to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Post(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("POST", url, a)
	return
}
func (req *Request) Post(url interface{}) (resp *Response, err error) {
	resp, err = Post(url2string(url), req2arg(req))
	return
}

// Put issues a PUT to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Put(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("PUT", url, a)
	return
}
func (req *Request) Put(url interface{}) (resp *Response, err error) {
	resp, err = Put(url2string(url), req2arg(req))
	return
}

// Patch issues a PATCH to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Patch(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("PATCH", url, a)
	return
}
func (req *Request) Patch(url interface{}) (resp *Response, err error) {
	resp, err = Patch(url2string(url), req2arg(req))
	return
}

// Delete issues a DELETE to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Delete(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("DELETE", url, a)
	return
}
func (req *Request) Delete(url interface{}) (resp *Response, err error) {
	resp, err = Delete(url2string(url), req2arg(req))
	return
}

// Options issues a OPTIONS to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Options(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("OPTIONS", url, a)
	return
}
func (req *Request) Options(url interface{}) (resp *Response, err error) {
	resp, err = Options(url2string(url), req2arg(req))
	return
}

func url2string(u interface{}) string {
	switch u.(type) {
	case string:
		return u.(string)
	case url.URL:
		s := u.(url.URL)
		return s.String()
	case *url.URL:
		s := u.(*url.URL)
		return s.String()
	}
	return ""
}

func req2arg(req *Request) (a *Args) {
	return &Args{
		Client:    req.Client,
		Headers:   req.Headers,
		Cookies:   req.Cookies,
		Data:      req.Data,
		Params:    req.Params,
		Files:     req.Files,
		Json:      req.Json,
		Proxy:     req.Proxy,
		BasicAuth: req.BasicAuth,
	}
}
<commit_msg>update comment<commit_after>package request

import (
	"io"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"

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

const Version = "0.2.0"

type FileField struct {
	FieldName string
	FileName  string
	File      io.Reader
}

type BasicAuth struct {
	Username string
	Password string
}

type Args struct {
	Client    *http.Client
	Headers   map[string]string
	Cookies   map[string]string
	Data      map[string]string
	Params    map[string]string
	Files     []FileField
	Json      interface{}
	Proxy     string
	BasicAuth BasicAuth
}

type Request struct {
	*Args
}

func NewArgs(c *http.Client) *Args {
	if c.Jar == nil {
		options := cookiejar.Options{
			PublicSuffixList: publicsuffix.List,
		}
		jar, _ := cookiejar.New(&options)
		c.Jar = jar
	}

	return &Args{
		Client:    c,
		Headers:   defaultHeaders,
		Cookies:   nil,
		Data:      nil,
		Params:    nil,
		Files:     nil,
		Json:      nil,
		Proxy:     "",
		BasicAuth: BasicAuth{},
	}
}

func NewRequest(c *http.Client) *Request {
	return &Request{NewArgs(c)}
}

func newURL(u string, params map[string]string) string {
	if params == nil {
		return u
	}

	p := url.Values{}
	for k, v := range params {
		p.Set(k, v)
	}
	if strings.Contains(u, "?") {
		return u + "&" + p.Encode()
	}
	return u + "?" + p.Encode()
}

func newBody(a *Args) (body io.Reader, contentType string, err error) {
	if a.Data == nil && a.Files == nil && a.Json == nil {
		return nil, "", nil
	}
	if a.Files != nil {
		return newMultipartBody(a)
	} else if a.Json != nil {
		return newJsonBody(a)
	}

	d := url.Values{}
	for k, v := range a.Data {
		d.Set(k, v)
	}
	return strings.NewReader(d.Encode()), "", nil
}

func newRequest(method string, url string, a *Args) (resp *Response, err error) {
	body, contentType, err := newBody(a)
	u := newURL(url, a.Params)
	req, err := http.NewRequest(method, u, body)
	if err != nil {
		return nil, err
	}
	applyHeaders(a, req, contentType)
	applyCookies(a, req)
	applyProxy(a)
	applyCheckRdirect(a)

	if a.BasicAuth.Username != "" {
		req.SetBasicAuth(a.BasicAuth.Username, a.BasicAuth.Password)
	}

	s, err := a.Client.Do(req)
	resp = &Response{s, nil}
	return
}

// Get issues a GET to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Get(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("GET", url, a)
	return
}

// url can be string or *url.URL
func (req *Request) Get(url interface{}) (resp *Response, err error) {
	resp, err = Get(url2string(url), req2arg(req))
	return
}

// Head issues a HEAD to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Head(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("HEAD", url, a)
	return
}

func (req *Request) Head(url interface{}) (resp *Response, err error) {
	resp, err = Head(url2string(url), req2arg(req))
	return
}

// Post issues a POST to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Post(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("POST", url, a)
	return
}
func (req *Request) Post(url interface{}) (resp *Response, err error) {
	resp, err = Post(url2string(url), req2arg(req))
	return
}

// Put issues a PUT to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Put(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("PUT", url, a)
	return
}
func (req *Request) Put(url interface{}) (resp *Response, err error) {
	resp, err = Put(url2string(url), req2arg(req))
	return
}

// Patch issues a PATCH to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Patch(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("PATCH", url, a)
	return
}
func (req *Request) Patch(url interface{}) (resp *Response, err error) {
	resp, err = Patch(url2string(url), req2arg(req))
	return
}

// Delete issues a DELETE to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Delete(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("DELETE", url, a)
	return
}
func (req *Request) Delete(url interface{}) (resp *Response, err error) {
	resp, err = Delete(url2string(url), req2arg(req))
	return
}

// Options issues a OPTIONS to the specified URL.
//
// Caller should close resp.Body when done reading from it.
func Options(url string, a *Args) (resp *Response, err error) {
	resp, err = newRequest("OPTIONS", url, a)
	return
}
func (req *Request) Options(url interface{}) (resp *Response, err error) {
	resp, err = Options(url2string(url), req2arg(req))
	return
}

func url2string(u interface{}) string {
	switch u.(type) {
	case string:
		return u.(string)
	case url.URL:
		s := u.(url.URL)
		return s.String()
	case *url.URL:
		s := u.(*url.URL)
		return s.String()
	}
	return ""
}

func req2arg(req *Request) (a *Args) {
	return &Args{
		Client:    req.Client,
		Headers:   req.Headers,
		Cookies:   req.Cookies,
		Data:      req.Data,
		Params:    req.Params,
		Files:     req.Files,
		Json:      req.Json,
		Proxy:     req.Proxy,
		BasicAuth: req.BasicAuth,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"log"

	"github.com/blang/semver"
	"github.com/rhysd/go-github-selfupdate/selfupdate"
	"github.com/urfave/cli"
)

func update(silent bool) error {
	if !silent {
		log.Printf("Checking for update")
	}

	v := semver.MustParse(version)
	updater, err := selfupdate.NewUpdater(selfupdate.Config{BinaryName: "rainforest"})
	if err != nil {
		return err
	}
	latest, err := updater.UpdateSelf(v, "rainforestapp/rainforest-cli")
	if err != nil {
		return err
	}
	if !silent {
		if latest.Version.Equals(v) {
			log.Println("No update available, already at the latest version!")
		} else {
			log.Printf("Updated to new version: %s!\n", latest.Version)
		}
	}

	return nil
}

func updateCmd(c cliContext) error {
	err := update(false)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}
	return nil
}

func autoUpdate(c cliContext, updateFinishedChan chan<- struct{}) {
	if !c.Bool("skip-update") {
		update(true)
	}
	updateFinishedChan <- struct{}{}
}
<commit_msg>RF-23055 Always output when CLI is updated<commit_after>package main

import (
	"log"

	"github.com/blang/semver"
	"github.com/rhysd/go-github-selfupdate/selfupdate"
	"github.com/urfave/cli"
)

func update(silent bool) error {
	if !silent {
		log.Printf("Checking for update")
	}

	v := semver.MustParse(version)
	updater, err := selfupdate.NewUpdater(selfupdate.Config{BinaryName: "rainforest"})
	if err != nil {
		return err
	}
	latest, err := updater.UpdateSelf(v, "rainforestapp/rainforest-cli")
	if err != nil {
		return err
	}
	if !silent && latest.Version.Equals(v) {
		log.Println("No update available, already at the latest version!")
	} else if latest.Version.NE(v) {
		log.Printf("Updated to new version: %s!\n", latest.Version)
	}

	return nil
}

func updateCmd(c cliContext) error {
	err := update(false)
	if err != nil {
		return cli.NewExitError(err.Error(), 1)
	}
	return nil
}

func autoUpdate(c cliContext, updateFinishedChan chan<- struct{}) {
	if !c.Bool("skip-update") {
		update(true)
	}
	updateFinishedChan <- struct{}{}
}
<|endoftext|>


<commit_before>package upsert

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"reflect"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
)

const (
	Inserted = iota
	Updated
	NoChange
)

func statusToText(status int) string {
	switch status {
	case Inserted:
		return "inserted"
	case Updated:
		return "updated"
	case NoChange:
		return "no change"
	}

	return "invalid status"
}

var (
	ErrNoIDReturned = errors.New("no id returned")

	// LongQuery will log long queries if set to a non-zero time
	LongQuery time.Duration
	Debug     = false
)

// Upserter is an interface specific to sqlx and PostgreSQL that can save a
// single row of data via Upsert(), Update() or Insert(). It doesn't try to
// know anything about relationships between tables. The behavior of Upserter
// depends on three struct tags.
//
//  * db: As with sqlx, this tag is the database column name for the field.
//     If db is not defined, the default is the lowercase value of the field
//     name.
//
//  * upsert: This may either be "key" or "omit". If it's "key", the
//     field/column is part of the where clause when attempting to update
//     an existing column. If it's "omit", the field is ignored completely.
//     By default, the field is considered a non-key value that should be
//     updated/set in the db.
//
//  * upsert_value: This is the placeholder for the value of the field for
//     use by sqlx.NamedExec(). By default, this is :column_name and typically
//     doesn't need to be changed. However, if the field needs to be
//     transformed by an SQL function before storing in the database,
//     this tag can be set. For example, if you had "lat" and "lon" columns,
//     you wouldn't want to store them in the db. Instead you'd want a
//     "location" column tagged with `upsert_value:"ll_to_earth(:lat, :lon)`
//
type Upserter interface {
	// Table returns table name we should save to
	Table() string
}

type columnSpec struct {
	name  string
	value string
}

func newColumnSpec(fieldName string, tag reflect.StructTag) columnSpec {
	cs := columnSpec{}

	// The name of the column is either the value of the "db" struct tag
	// or a lowercase version of the field name.
	dbTag := tag.Get("db")
	if len(dbTag) > 0 {
		cs.name = dbTag
	} else {
		cs.name = strings.ToLower(fieldName)
	}

	// The value placeholder of the column is typically just ":column_name"
	// but can be overriden with upsert_value.
	val := tag.Get("upsert_value")
	if len(val) > 0 {
		cs.value = val
	} else {
		cs.value = ":" + cs.name
	}

	return cs
}

// updateColumns returns the fields that are read from the struct and set
// on upserting in the db. Typically this should include everything except the
// key fields and any composite (array, nested struct) types or any
// field that doesn't map directly into a db column. Tag a field with
// `upsert:"omit"` to explicitly exclude from this list.
func updateColumns(u interface{}) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag

		// Include any column that isn't tagged with upsert:omit
		if !strings.Contains(tag.Get("upsert"), "omit") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// uniqueKeyColumns returns the fields of the struct that together are
// naturally unique. For example, an md5 hash of the content. Or a
// foreign key plus an internal value. This is used in where clause
// when trying to find existing rows. Tag a field with `"upsert:"key"`
// to include in the unique key.
func uniqueKeyColumns(u interface{}) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag
		// Check if upsert tag contains "key". This wouldn't work
		// if possible options were substrings of one another. For a
		// better implementation, look at src/encoding/json/tags.go
		if strings.Contains(tag.Get("upsert"), "key") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// set returns a string like "SET "col1" = :col1, "col2" = :col2" for
// use with sqlx.NamedExec() and friends.
func set(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteString("SET ")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s" = %s`, cols[i].name, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}

	return b.String()
}

// values returns a string like `("col1", "col2") VALUES(:col1, :col2)`
// for use with sqlx.NamedExec() etc.
func values(u interface{}) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteRune('(')
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s"`, cols[i].name)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	b.WriteString("VALUES (")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `%s`, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	return b.String()
}

// where returns an SQL where clause with all the key columns of
// this Upserter
func where(u Upserter) string {
	keycols := uniqueKeyColumns(u)
	b := bytes.Buffer{}
	n := len(keycols)

	fmt.Fprintf(&b, "WHERE ")

	for i := 0; i < n; i++ {
		// If we need to support NULLs here, the best option may be
		// something like (x = y OR (x is null and y is null))
		// rather than "IS NOT DISTINCT FROM" which doesn't use indexes
		// it seems
		fmt.Fprintf(&b, `%s = %s`, keycols[i].name, keycols[i].value)

		if i < n-1 {
			fmt.Fprint(&b, " AND ")
		}
	}

	return b.String()
}

// updateSQL returns a full SQL command to update this Upserter u
func updateSQL(u Upserter) string {
	q := fmt.Sprintf(`UPDATE "%s" %s %s RETURNING *`,
		u.Table(), set(u), where(u))

	return q
}

// insertSQL returns a full SQL command to insert this Upserter u
func insertSQL(u Upserter) string {
	q := fmt.Sprintf(`INSERT INTO "%s" %s RETURNING *`, u.Table(), values(u))

	return q
}

// getSQL returns a full SQL command to retrieve this Upserter u
func getSQL(u Upserter) string {
	q := fmt.Sprintf(`SELECT * FROM %s %s`, u.Table(), where(u))

	return q
}

func Update(ext sqlx.Ext, u Upserter) (status int, err error) {
	q := updateSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	otherPtr := reflect.New(reflect.TypeOf(u).Elem())
	other := reflect.Indirect(otherPtr)
	otherInterface := other.Addr().Interface()

	rows, err := sqlx.NamedQuery(ext, getSQL(u), u)
	if err != nil {
		log.Println("error getting", err)
		return
	}

	if rows.Next() {
		err = rows.StructScan(otherInterface)
		if err != nil {
			log.Println("error scanning", err)
			return
		}

		if values(otherInterface) == values(u) {
			status = NoChange
			rows.Close()
			return
		}
	}
	rows.Close()

	status = Updated

	// Try to update an existing row
	rows, err = sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(updateSQL(u), err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// We could not find anything to update.
		err = ErrNoIDReturned
		return
	}

	return
}

// Insert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert it
// and set its primary key id value.
func Insert(ext sqlx.Ext, u Upserter) (err error) {
	q := insertSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	// Try to insert a row
	rows, err := sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// No rows were returned but no SQL error. Weird, return generic
		// error.
		err = ErrNoIDReturned
		return
	}

	return
}

// Upsert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert/update it
// and set the new primary key id if that succeeds. inserted returns true
// if a new row was inserted. The client is responsible for wrapping
// in a transaction when needed. This can be used when running a transaction
// at a higher level (upserting multiple items).
func Upsert(ext sqlx.Ext, u Upserter) (status int, err error) {
	defer func() {
		if Debug {
			log.Println(statusToText(status), u)
		}
	}()

	// Try to update, return immediately if succcesful
	status, err = Update(ext, u)
	if err == nil {
		return
	}

	// Can't update? Try insert
	err = Insert(ext, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

// UpsertTx takes only an sqlx.DB and wraps the upsert attempt into a
// a transaction.
func UpsertTx(db *sqlx.DB, u Upserter) (status int, err error) {
	defer func() {
		if Debug {
			log.Println(statusToText(status), u)
		}
	}()

	tx, err := db.Beginx()
	if err != nil {
		log.Println("can't start transaction", err)
		return
	}
	defer func() {
		if err != nil {
			tx.Rollback()
		} else {
			tx.Commit()
		}
	}()

	// Try to update
	status, err = Update(tx, u)

	// If we have a nil error, we successfully updated. If we have
	// an err other than ErrNoIDReturned, we couldn't update for an
	// unexpected reason. In either case return.
	if err != ErrNoIDReturned {
		return
	}

	// No ID returned in the update? Try insert
	err = Insert(tx, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	if Debug {
		log.Println(statusToText(status), u)
	}

	return
}

func Delete(ext sqlx.Ext, u Upserter) (err error) {
	q := fmt.Sprintf(`DELETE FROM "%s" %s`,
		u.Table(), where(u))
	_, err = sqlx.NamedExec(ext, q, u)

	if err != nil {
		log.Println("can't delete", err)
		return
	}

	return
}
<commit_msg>also use values to fix test<commit_after>package upsert

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"reflect"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
)

const (
	Inserted = iota
	Updated
	NoChange
)

func statusToText(status int) string {
	switch status {
	case Inserted:
		return "inserted"
	case Updated:
		return "updated"
	case NoChange:
		return "no change"
	}

	return "invalid status"
}

var (
	ErrNoIDReturned = errors.New("no id returned")

	// LongQuery will log long queries if set to a non-zero time
	LongQuery time.Duration
	Debug     = false
)

// Upserter is an interface specific to sqlx and PostgreSQL that can save a
// single row of data via Upsert(), Update() or Insert(). It doesn't try to
// know anything about relationships between tables. The behavior of Upserter
// depends on three struct tags.
//
//  * db: As with sqlx, this tag is the database column name for the field.
//     If db is not defined, the default is the lowercase value of the field
//     name.
//
//  * upsert: This may either be "key" or "omit". If it's "key", the
//     field/column is part of the where clause when attempting to update
//     an existing column. If it's "omit", the field is ignored completely.
//     By default, the field is considered a non-key value that should be
//     updated/set in the db.
//
//  * upsert_value: This is the placeholder for the value of the field for
//     use by sqlx.NamedExec(). By default, this is :column_name and typically
//     doesn't need to be changed. However, if the field needs to be
//     transformed by an SQL function before storing in the database,
//     this tag can be set. For example, if you had "lat" and "lon" columns,
//     you wouldn't want to store them in the db. Instead you'd want a
//     "location" column tagged with `upsert_value:"ll_to_earth(:lat, :lon)`
//
type Upserter interface {
	// Table returns table name we should save to
	Table() string
}

type columnSpec struct {
	name  string
	value string
}

func newColumnSpec(fieldName string, tag reflect.StructTag) columnSpec {
	cs := columnSpec{}

	// The name of the column is either the value of the "db" struct tag
	// or a lowercase version of the field name.
	dbTag := tag.Get("db")
	if len(dbTag) > 0 {
		cs.name = dbTag
	} else {
		cs.name = strings.ToLower(fieldName)
	}

	// The value placeholder of the column is typically just ":column_name"
	// but can be overriden with upsert_value.
	val := tag.Get("upsert_value")
	if len(val) > 0 {
		cs.value = val
	} else {
		cs.value = ":" + cs.name
	}

	return cs
}

// updateColumns returns the fields that are read from the struct and set
// on upserting in the db. Typically this should include everything except the
// key fields and any composite (array, nested struct) types or any
// field that doesn't map directly into a db column. Tag a field with
// `upsert:"omit"` to explicitly exclude from this list.
func updateColumns(u interface{}) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag

		// Include any column that isn't tagged with upsert:omit
		if !strings.Contains(tag.Get("upsert"), "omit") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

func updateValueString(u interface{}) string {
	b := bytes.NewBuffer(nil)

	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return ""
	}

	val := reflect.ValueOf(u).Elem()

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag

		// Include any column that isn't tagged with upsert:omit
		if !strings.Contains(tag.Get("upsert"), "omit") {
			fmt.Fprintf(b, "%v ", val.Field(i))
		}
	}

	return b.String()
}

// uniqueKeyColumns returns the fields of the struct that together are
// naturally unique. For example, an md5 hash of the content. Or a
// foreign key plus an internal value. This is used in where clause
// when trying to find existing rows. Tag a field with `"upsert:"key"`
// to include in the unique key.
func uniqueKeyColumns(u interface{}) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag
		// Check if upsert tag contains "key". This wouldn't work
		// if possible options were substrings of one another. For a
		// better implementation, look at src/encoding/json/tags.go
		if strings.Contains(tag.Get("upsert"), "key") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// set returns a string like "SET "col1" = :col1, "col2" = :col2" for
// use with sqlx.NamedExec() and friends.
func set(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteString("SET ")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s" = %s`, cols[i].name, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}

	return b.String()
}

// values returns a string like `("col1", "col2") VALUES(:col1, :col2)`
// for use with sqlx.NamedExec() etc.
func values(u interface{}) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteRune('(')
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s"`, cols[i].name)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	b.WriteString("VALUES (")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `%s`, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	return b.String()
}

// where returns an SQL where clause with all the key columns of
// this Upserter
func where(u Upserter) string {
	keycols := uniqueKeyColumns(u)
	b := bytes.Buffer{}
	n := len(keycols)

	fmt.Fprintf(&b, "WHERE ")

	for i := 0; i < n; i++ {
		// If we need to support NULLs here, the best option may be
		// something like (x = y OR (x is null and y is null))
		// rather than "IS NOT DISTINCT FROM" which doesn't use indexes
		// it seems
		fmt.Fprintf(&b, `%s = %s`, keycols[i].name, keycols[i].value)

		if i < n-1 {
			fmt.Fprint(&b, " AND ")
		}
	}

	return b.String()
}

// updateSQL returns a full SQL command to update this Upserter u
func updateSQL(u Upserter) string {
	q := fmt.Sprintf(`UPDATE "%s" %s %s RETURNING *`,
		u.Table(), set(u), where(u))

	return q
}

// insertSQL returns a full SQL command to insert this Upserter u
func insertSQL(u Upserter) string {
	q := fmt.Sprintf(`INSERT INTO "%s" %s RETURNING *`, u.Table(), values(u))

	return q
}

// getSQL returns a full SQL command to retrieve this Upserter u
func getSQL(u Upserter) string {
	q := fmt.Sprintf(`SELECT * FROM %s %s`, u.Table(), where(u))

	return q
}

func Update(ext sqlx.Ext, u Upserter) (status int, err error) {
	q := updateSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	otherPtr := reflect.New(reflect.TypeOf(u).Elem())
	other := reflect.Indirect(otherPtr)
	otherInterface := other.Addr().Interface()

	rows, err := sqlx.NamedQuery(ext, getSQL(u), u)
	if err != nil {
		log.Println("error getting", err)
		return
	}

	if rows.Next() {
		err = rows.StructScan(otherInterface)
		if err != nil {
			log.Println("error scanning", err)
			return
		}

		if values(otherInterface) == values(u) && updateValueString(otherInterface) == updateValueString(u) {
			status = NoChange
			rows.Close()
			return
		}
	}
	rows.Close()

	status = Updated

	// Try to update an existing row
	rows, err = sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(updateSQL(u), err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// We could not find anything to update.
		err = ErrNoIDReturned
		return
	}

	return
}

// Insert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert it
// and set its primary key id value.
func Insert(ext sqlx.Ext, u Upserter) (err error) {
	q := insertSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	// Try to insert a row
	rows, err := sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// No rows were returned but no SQL error. Weird, return generic
		// error.
		err = ErrNoIDReturned
		return
	}

	return
}

// Upsert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert/update it
// and set the new primary key id if that succeeds. inserted returns true
// if a new row was inserted. The client is responsible for wrapping
// in a transaction when needed. This can be used when running a transaction
// at a higher level (upserting multiple items).
func Upsert(ext sqlx.Ext, u Upserter) (status int, err error) {
	defer func() {
		if Debug {
			log.Println(statusToText(status), u)
		}
	}()

	// Try to update, return immediately if succcesful
	status, err = Update(ext, u)
	if err == nil {
		return
	}

	// Can't update? Try insert
	err = Insert(ext, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

// UpsertTx takes only an sqlx.DB and wraps the upsert attempt into a
// a transaction.
func UpsertTx(db *sqlx.DB, u Upserter) (status int, err error) {
	defer func() {
		if Debug {
			log.Println(statusToText(status), u)
		}
	}()

	tx, err := db.Beginx()
	if err != nil {
		log.Println("can't start transaction", err)
		return
	}
	defer func() {
		if err != nil {
			tx.Rollback()
		} else {
			tx.Commit()
		}
	}()

	// Try to update
	status, err = Update(tx, u)

	// If we have a nil error, we successfully updated. If we have
	// an err other than ErrNoIDReturned, we couldn't update for an
	// unexpected reason. In either case return.
	if err != ErrNoIDReturned {
		return
	}

	// No ID returned in the update? Try insert
	err = Insert(tx, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	if Debug {
		log.Println(statusToText(status), u)
	}

	return
}

func Delete(ext sqlx.Ext, u Upserter) (err error) {
	q := fmt.Sprintf(`DELETE FROM "%s" %s`,
		u.Table(), where(u))
	_, err = sqlx.NamedExec(ext, q, u)

	if err != nil {
		log.Println("can't delete", err)
		return
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"
	"os"
)

var (
	silent bool
	// Declared in common:
	// hashFunction HashValue
	// cwd string
)

var cmdVerify = &Command{
	Run: verify,
	Usage: `verify [-s] [-c=hash] [-D=dir]`,
	Short: "Verify using a checksum file",
	Long: `
Verify all files for the given directory against a checksum file.`,
}

// Initialized in common:
// var preferredHashes []string

func init() {
	const (
		silentUsage = "silent; don't output to STDERR"
		cwdUsage = "the directory for which to verify using the checksum file"
	)
	hashUsage := fmt.Sprintf("the hash to use; if unspecified, the following are tried in order: %s", strings.Join(preferredHashes, ", "))
	f := &cmdVerify.Flag
	f.BoolVar(&silent, "s", false, silentUsage)
	f.Var(&hashFunction, "c", hashUsage)
	f.StringVar(&cwd, "D", ".", cwdUsage)
}

func verify(cmd *Command, args []string) {
	var err error
	hash, checksums := setDirAndHashOptions()

	allMatch := true
	r := NewReader(checksums)
	err = r.Each(func(entry *Entry) {
		if ok, err := entry.Verify(hash.Hash); err != nil {
			if !silent {
				warn("%s\n", err)
			}
			allMatch = false
		} else if !ok {
			if !silent {
				warn("%s does not match %s\n", entry.Filename, entry.Checksum)
			}
			allMatch = false
		}
	})
	if err != nil {
		die(err)
	}

	if allMatch {
		os.Exit(0)
	} else {
		os.Exit(1)
	}
}

<commit_msg>Add a single file mode for the verify command<commit_after>package main

import (
	"fmt"
	"strings"
	"os"
)

var (
	silent bool
	// Declared in common:
	// hashFunction HashValue
	// cwd string
)

var cmdVerify = &Command{
	Run: verify,
	Usage: `verify [-s] [-c=hash] [-D=dir] [file]`,
	Short: "Verify using a checksum file",
	Long: `
Verify all files for the given directory against a checksum file. If a specific
file is given, only that file is verified.`,
}

// Initialized in common:
// var preferredHashes []string

func init() {
	const (
		silentUsage = "silent; don't output to STDERR"
		cwdUsage = "the directory for which to verify using the checksum file"
	)
	hashUsage := fmt.Sprintf("the hash to use; if unspecified, the following are tried in order: %s", strings.Join(preferredHashes, ", "))
	f := &cmdVerify.Flag
	f.BoolVar(&silent, "s", false, silentUsage)
	f.Var(&hashFunction, "c", hashUsage)
	f.StringVar(&cwd, "D", ".", cwdUsage)
}

func verify(cmd *Command, args []string) {
	var err error
	hash, checksums := setDirAndHashOptions()
	singleFileMode, singleFile := false, ""
	if len(args) == 1 {
		singleFileMode = true
		singleFile = args[0]
	}

	allMatch := true
	r := NewReader(checksums)
	err = r.Each(func(entry *Entry) {
		if singleFileMode && entry.Filename != singleFile {
			return // continue
		}

		if ok, err := entry.Verify(hash.Hash); err != nil {
			if !silent {
				warn("%s\n", err)
			}
			allMatch = false
		} else if !ok {
			if !silent {
				warn("%s does not match %s\n", entry.Filename, entry.Checksum)
			}
			allMatch = false
		}
	})
	if err != nil {
		die(err)
	}

	if allMatch {
		os.Exit(0)
	} else {
		os.Exit(1)
	}
}

<|endoftext|>


<commit_before>package containercommands

import (
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/jrperritt/rack/handler"
	"github.com/jrperritt/rack/internal/github.com/cenkalti/backoff"
	"github.com/jrperritt/rack/internal/github.com/dustin/go-humanize"
	"github.com/jrperritt/rack/internal/github.com/rackspace/gophercloud/openstack/objectstorage/v1/objects"
	"github.com/jrperritt/rack/internal/github.com/rackspace/gophercloud/rackspace/objectstorage/v1/containers"
	"github.com/jrperritt/rack/util"
)

type handleEmptyParams struct {
	container   string
	quiet       bool
	concurrency int
}

func handleEmpty(command handler.Commander, resource *handler.Resource, params *handleEmptyParams) {
	var totalFiles int64
	var wg sync.WaitGroup

	// bump thread count to number of available CPUs
	runtime.GOMAXPROCS(runtime.NumCPU())

	// get the names of all the objects in the container
	allPages, err := objects.List(command.Context().ServiceClient, params.container, nil).AllPages()
	if err != nil {
		resource.Err = err
		return
	}
	names, err := objects.ExtractNames(allPages)
	if err != nil {
		resource.Err = err
		return
	}

	// send the object names into the `jobs` channel
	jobs := make(chan string, len(names))
	for i := 0; i < len(names); i++ {
		wg.Add(1)
		jobs <- names[i]
	}

	// default the number of goroutines to spawn if the `concurrency` flag
	// wasn't provided
	if params.concurrency == 0 {
		params.concurrency = 100
	}

	start := time.Now()

	for i := 0; i < params.concurrency; i++ {
		go func(totalFiles *int64) {
			for objectName := range jobs {
				ticker := backoff.NewTicker(backoff.NewExponentialBackOff())
				for _ = range ticker.C {
					rawResponse := objects.Delete(command.Context().ServiceClient, params.container, objectName, nil)
					if rawResponse.Err != nil {
						jobs <- objectName
						continue
					}
					ticker.Stop()
					break
				}

				*totalFiles++

				if !params.quiet {
					re := &handler.Resource{
						Result: fmt.Sprintf("Successfully deleted object [%s] from container [%s]\n", objectName, params.container),
					}
					command.Context().Results <- re
				}
				wg.Done()
			}
		}(&totalFiles)
	}

	fmt.Println("Waiting for wait group to be done")
	wg.Wait()

	getResult, err := containers.Get(command.Context().ServiceClient, params.container).Extract()
	fmt.Printf("getResult: %+v\n", getResult)
	if err == nil && getResult.ObjectCount != 0 {
		fmt.Printf("Re-running handleEmpty because not all objects were deleted. Still have %d\n", getResult.ObjectCount)
		params.concurrency = 1
		handleEmpty(command, resource, params)
	}

	resource.Result = fmt.Sprintf("Finished! Deleted %s %s in %s", humanize.Comma(totalFiles), util.Pluralize("object", totalFiles), humanize.RelTime(start, time.Now(), "", ""))
}
<commit_msg>remove debugging code from handleEmpty<commit_after>package containercommands

import (
	"fmt"
	"runtime"
	"sync"
	"time"

	"github.com/jrperritt/rack/handler"
	"github.com/jrperritt/rack/internal/github.com/cenkalti/backoff"
	"github.com/jrperritt/rack/internal/github.com/dustin/go-humanize"
	"github.com/jrperritt/rack/internal/github.com/rackspace/gophercloud/openstack/objectstorage/v1/objects"
	"github.com/jrperritt/rack/util"
)

type handleEmptyParams struct {
	container   string
	quiet       bool
	concurrency int
}

func handleEmpty(command handler.Commander, resource *handler.Resource, params *handleEmptyParams) {
	var totalFiles int64
	var wg sync.WaitGroup

	// bump thread count to number of available CPUs
	runtime.GOMAXPROCS(runtime.NumCPU())

	// get the names of all the objects in the container
	allPages, err := objects.List(command.Context().ServiceClient, params.container, nil).AllPages()
	if err != nil {
		resource.Err = err
		return
	}
	names, err := objects.ExtractNames(allPages)
	if err != nil {
		resource.Err = err
		return
	}

	// send the object names into the `jobs` channel
	jobs := make(chan string, len(names))
	for i := 0; i < len(names); i++ {
		wg.Add(1)
		jobs <- names[i]
	}
	close(jobs)

	// default the number of goroutines to spawn if the `concurrency` flag
	// wasn't provided
	if params.concurrency == 0 {
		params.concurrency = 100
	}

	start := time.Now()

	for i := 0; i < params.concurrency; i++ {
		go func(totalFiles *int64) {
			for objectName := range jobs {
				ticker := backoff.NewTicker(backoff.NewExponentialBackOff())
				for _ = range ticker.C {
					rawResponse := objects.Delete(command.Context().ServiceClient, params.container, objectName, nil)
					if rawResponse.Err != nil {
						continue
					}
					ticker.Stop()
					break
				}

				*totalFiles++

				if !params.quiet {
					re := &handler.Resource{
						Result: fmt.Sprintf("Successfully deleted object [%s] from container [%s]\n", objectName, params.container),
					}
					command.Context().Results <- re
				}
				wg.Done()
			}
		}(&totalFiles)
	}

	wg.Wait()

	resource.Result = fmt.Sprintf("Finished! Deleted %s %s in %s", humanize.Comma(totalFiles), util.Pluralize("object", totalFiles), humanize.RelTime(start, time.Now(), "", ""))
}
<|endoftext|>


<commit_before>package devicepathresolver_test

import (
	"runtime"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"errors"

	. "github.com/cloudfoundry/bosh-agent/infrastructure/devicepathresolver"
	boshsettings "github.com/cloudfoundry/bosh-agent/settings"
	fakesys "github.com/cloudfoundry/bosh-utils/system/fakes"
)

var _ = Describe("mappedDevicePathResolver", func() {
	var (
		fs           *fakesys.FakeFileSystem
		diskSettings boshsettings.DiskSettings
		resolver     DevicePathResolver
	)

	BeforeEach(func() {
		if runtime.GOOS == "windows" {
			Skip("Not yet implemented on Windows")
		}

		fs = fakesys.NewFakeFileSystem()
		resolver = NewMappedDevicePathResolver(time.Second, fs)
		diskSettings = boshsettings.DiskSettings{
			Path: "/dev/sda",
		}
	})

	Context("when path is not provided", func() {
		It("returns an error", func() {
			diskSettings := boshsettings.DiskSettings{}
			_, _, err := resolver.GetRealDevicePath(diskSettings)
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("path is missing"))
		})
	})

	Context("when a matching /dev/xvdX device is found", func() {
		BeforeEach(func() {
			fs.WriteFile("/dev/xvda", []byte{})
			fs.WriteFile("/dev/vda", []byte{})
			fs.WriteFile("/dev/sda", []byte{})
		})

		It("returns the match", func() {
			realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
			Expect(err).NotTo(HaveOccurred())
			Expect(timedOut).To(BeFalse())
			Expect(realPath).To(Equal("/dev/xvda"))
		})
	})

	Context("when a matching /dev/vdX device is found", func() {
		BeforeEach(func() {
			fs.WriteFile("/dev/vda", []byte{})
			fs.WriteFile("/dev/sda", []byte{})
		})

		It("returns the match", func() {
			realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
			Expect(err).NotTo(HaveOccurred())
			Expect(timedOut).To(BeFalse())
			Expect(realPath).To(Equal("/dev/vda"))
		})
	})

	Context("when a matching /dev/sdX device is found", func() {
		BeforeEach(func() {
			fs.WriteFile("/dev/sda", []byte{})
		})

		It("returns the match", func() {
			realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
			Expect(err).NotTo(HaveOccurred())
			Expect(timedOut).To(BeFalse())
			Expect(realPath).To(Equal("/dev/sda"))
		})
	})

	Context("when no matching device is found the first time", func() {
		Context("when the timeout has not expired", func() {
			BeforeEach(func() {
				time.AfterFunc(time.Second, func() {
					fs.WriteFile("/dev/xvda", []byte{})
				})
			})

			It("returns the match", func() {
				realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
				Expect(err).NotTo(HaveOccurred())
				Expect(timedOut).To(BeFalse())
				Expect(realPath).To(Equal("/dev/xvda"))
			})
		})

		Context("when the timeout has expired", func() {
			It("errs", func() {
				_, timedOut, err := resolver.GetRealDevicePath(diskSettings)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("Timed out getting real device path for /dev/sda"))
				Expect(timedOut).To(BeTrue())
			})
		})
	})

	Context("when a path that never needs remapping is passed in", func() {
		Context("when path exists", func() {
			BeforeEach(func() {
				fs.WriteFile("/dev/xvdba", []byte{})
				diskSettings = boshsettings.DiskSettings{
					Path: "/dev/xvdba",
				}
			})

			It("returns the path as given", func() {
				realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
				Expect(realPath).To(Equal("/dev/xvdba"))
				Expect(err).ToNot(HaveOccurred())
				Expect(timedOut).To(BeFalse())
			})
		})

		FContext("when path is a symlink", func() {
			BeforeEach(func() {
				Expect(fs.WriteFile("/dev/xvdba", []byte{})).To(Succeed())
				Expect(fs.MkdirAll("/dev/disk/by-label", 0755)).To(Succeed())

				diskSettings = boshsettings.DiskSettings{
					Path: "/dev/disk/by-label/my-disk-label",
				}
			})

			Context("and resolves to absolute path", func() {
				BeforeEach(func() {
					Expect(fs.Symlink("/dev/xvdba", "/dev/disk/by-label/my-disk-label")).To(Succeed())
				})

				It("returns the resolved path", func() {
					realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
					Expect(realPath).To(Equal("/dev/xvdba"))
					Expect(err).ToNot(HaveOccurred())
					Expect(timedOut).To(BeFalse())
				})
			})

			Context("and resolves to a relative path", func() {
				BeforeEach(func() {
					Expect(fs.Symlink("../../xvdba", "/dev/disk/by-label/my-disk-label")).To(Succeed())
				})

				It("returns the resolved path", func() {
					realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
					Expect(realPath).To(Equal("/dev/xvdba"))
					Expect(err).ToNot(HaveOccurred())
					Expect(timedOut).To(BeFalse())
				})
			})

			Context("and ReadLink fails", func() {
				BeforeEach(func() {
					Expect(fs.Symlink("/dev/xvdba", "/dev/disk/by-label/my-disk-label")).To(Succeed())
					fs.ReadlinkError = errors.New("can't read symlink")
				})

				It("returns error", func() {
					_, timedOut, err := resolver.GetRealDevicePath(diskSettings)
					Expect(err).To(HaveOccurred())
					Expect(err.Error()).To(Equal("Getting real device path: can't read symlink"))
					Expect(timedOut).To(BeFalse())
				})
			})
		})

		Context("when path does not exist", func() {
			BeforeEach(func() {
				diskSettings = boshsettings.DiskSettings{
					Path: "/dev/xvdba",
				}
			})

			It("returns an error", func() {
				realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
				Expect(realPath).To(Equal(""))
				Expect(err).To(HaveOccurred())
				Expect(timedOut).To(BeTrue())
			})
		})
	})
})
<commit_msg>Remove focus<commit_after>package devicepathresolver_test

import (
	"runtime"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"errors"

	. "github.com/cloudfoundry/bosh-agent/infrastructure/devicepathresolver"
	boshsettings "github.com/cloudfoundry/bosh-agent/settings"
	fakesys "github.com/cloudfoundry/bosh-utils/system/fakes"
)

var _ = Describe("mappedDevicePathResolver", func() {
	var (
		fs           *fakesys.FakeFileSystem
		diskSettings boshsettings.DiskSettings
		resolver     DevicePathResolver
	)

	BeforeEach(func() {
		if runtime.GOOS == "windows" {
			Skip("Not yet implemented on Windows")
		}

		fs = fakesys.NewFakeFileSystem()
		resolver = NewMappedDevicePathResolver(time.Second, fs)
		diskSettings = boshsettings.DiskSettings{
			Path: "/dev/sda",
		}
	})

	Context("when path is not provided", func() {
		It("returns an error", func() {
			diskSettings := boshsettings.DiskSettings{}
			_, _, err := resolver.GetRealDevicePath(diskSettings)
			Expect(err).To(HaveOccurred())
			Expect(err.Error()).To(ContainSubstring("path is missing"))
		})
	})

	Context("when a matching /dev/xvdX device is found", func() {
		BeforeEach(func() {
			fs.WriteFile("/dev/xvda", []byte{})
			fs.WriteFile("/dev/vda", []byte{})
			fs.WriteFile("/dev/sda", []byte{})
		})

		It("returns the match", func() {
			realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
			Expect(err).NotTo(HaveOccurred())
			Expect(timedOut).To(BeFalse())
			Expect(realPath).To(Equal("/dev/xvda"))
		})
	})

	Context("when a matching /dev/vdX device is found", func() {
		BeforeEach(func() {
			fs.WriteFile("/dev/vda", []byte{})
			fs.WriteFile("/dev/sda", []byte{})
		})

		It("returns the match", func() {
			realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
			Expect(err).NotTo(HaveOccurred())
			Expect(timedOut).To(BeFalse())
			Expect(realPath).To(Equal("/dev/vda"))
		})
	})

	Context("when a matching /dev/sdX device is found", func() {
		BeforeEach(func() {
			fs.WriteFile("/dev/sda", []byte{})
		})

		It("returns the match", func() {
			realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
			Expect(err).NotTo(HaveOccurred())
			Expect(timedOut).To(BeFalse())
			Expect(realPath).To(Equal("/dev/sda"))
		})
	})

	Context("when no matching device is found the first time", func() {
		Context("when the timeout has not expired", func() {
			BeforeEach(func() {
				time.AfterFunc(time.Second, func() {
					fs.WriteFile("/dev/xvda", []byte{})
				})
			})

			It("returns the match", func() {
				realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
				Expect(err).NotTo(HaveOccurred())
				Expect(timedOut).To(BeFalse())
				Expect(realPath).To(Equal("/dev/xvda"))
			})
		})

		Context("when the timeout has expired", func() {
			It("errs", func() {
				_, timedOut, err := resolver.GetRealDevicePath(diskSettings)
				Expect(err).To(HaveOccurred())
				Expect(err.Error()).To(Equal("Timed out getting real device path for /dev/sda"))
				Expect(timedOut).To(BeTrue())
			})
		})
	})

	Context("when a path that never needs remapping is passed in", func() {
		Context("when path exists", func() {
			BeforeEach(func() {
				fs.WriteFile("/dev/xvdba", []byte{})
				diskSettings = boshsettings.DiskSettings{
					Path: "/dev/xvdba",
				}
			})

			It("returns the path as given", func() {
				realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
				Expect(realPath).To(Equal("/dev/xvdba"))
				Expect(err).ToNot(HaveOccurred())
				Expect(timedOut).To(BeFalse())
			})
		})

		Context("when path is a symlink", func() {
			BeforeEach(func() {
				Expect(fs.WriteFile("/dev/xvdba", []byte{})).To(Succeed())
				Expect(fs.MkdirAll("/dev/disk/by-label", 0755)).To(Succeed())

				diskSettings = boshsettings.DiskSettings{
					Path: "/dev/disk/by-label/my-disk-label",
				}
			})

			Context("and resolves to absolute path", func() {
				BeforeEach(func() {
					Expect(fs.Symlink("/dev/xvdba", "/dev/disk/by-label/my-disk-label")).To(Succeed())
				})

				It("returns the resolved path", func() {
					realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
					Expect(realPath).To(Equal("/dev/xvdba"))
					Expect(err).ToNot(HaveOccurred())
					Expect(timedOut).To(BeFalse())
				})
			})

			Context("and resolves to a relative path", func() {
				BeforeEach(func() {
					Expect(fs.Symlink("../../xvdba", "/dev/disk/by-label/my-disk-label")).To(Succeed())
				})

				It("returns the resolved path", func() {
					realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
					Expect(realPath).To(Equal("/dev/xvdba"))
					Expect(err).ToNot(HaveOccurred())
					Expect(timedOut).To(BeFalse())
				})
			})

			Context("and ReadLink fails", func() {
				BeforeEach(func() {
					Expect(fs.Symlink("/dev/xvdba", "/dev/disk/by-label/my-disk-label")).To(Succeed())
					fs.ReadlinkError = errors.New("can't read symlink")
				})

				It("returns error", func() {
					_, timedOut, err := resolver.GetRealDevicePath(diskSettings)
					Expect(err).To(HaveOccurred())
					Expect(err.Error()).To(Equal("Getting real device path: can't read symlink"))
					Expect(timedOut).To(BeFalse())
				})
			})
		})

		Context("when path does not exist", func() {
			BeforeEach(func() {
				diskSettings = boshsettings.DiskSettings{
					Path: "/dev/xvdba",
				}
			})

			It("returns an error", func() {
				realPath, timedOut, err := resolver.GetRealDevicePath(diskSettings)
				Expect(realPath).To(Equal(""))
				Expect(err).To(HaveOccurred())
				Expect(timedOut).To(BeTrue())
			})
		})
	})
})
<|endoftext|>


<commit_before>// +build linux

package systemd

import (
	"bytes"
	"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/godbus/dbus"
)

type systemdCgroup struct {
	cgroup *cgroups.Cgroup
}

type subsystem interface {
	GetStats(string, *cgroups.Stats) error
}

var (
	connLock              sync.Mutex
	theConn               *systemd.Conn
	hasStartTransientUnit 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
}

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

func Apply(c *cgroups.Cgroup, pid int) (map[string]string, error) {
	var (
		unitName   = getUnitName(c)
		slice      = "system.slice"
		properties []systemd.Property
		res        = &systemdCgroup{}
	)

	res.cgroup = c

	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 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 nil, err
	}

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

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

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

	paths := make(map[string]string)
	for _, sysname := range []string{
		"devices",
		"memory",
		"cpu",
		"cpuset",
		"cpuacct",
		"blkio",
		"perf_event",
		"freezer",
	} {
		subsystemPath, err := getSubsystemPath(res.cgroup, sysname)
		if err != nil {
			// Don't fail if a cgroup hierarchy was not found, just skip this subsystem
			if cgroups.IsNotFound(err) {
				continue
			}
			return nil, err
		}
		paths[sysname] = subsystemPath
	}
	return paths, nil
}

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

func joinFreezer(c *cgroups.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "freezer")
	if err != nil {
		return err
	}

	if err := os.MkdirAll(path, 0755); err != nil && !os.IsExist(err) {
		return err
	}

	return ioutil.WriteFile(filepath.Join(path, "cgroup.procs"), []byte(strconv.Itoa(pid)), 0700)
}

func getSubsystemPath(c *cgroups.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 Freeze(c *cgroups.Cgroup, state cgroups.FreezerState) error {
	path, err := getSubsystemPath(c, "freezer")
	if err != nil {
		return err
	}

	if err := ioutil.WriteFile(filepath.Join(path, "freezer.state"), []byte(state), 0); err != nil {
		return err
	}
	for {
		state_, err := ioutil.ReadFile(filepath.Join(path, "freezer.state"))
		if err != nil {
			return err
		}
		if string(state) == string(bytes.TrimSpace(state_)) {
			break
		}
		time.Sleep(1 * time.Millisecond)
	}
	return nil
}

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

	return cgroups.ReadProcsFile(path)
}

func getUnitName(c *cgroups.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 *cgroups.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "devices")
	if err != nil {
		return err
	}

	if err := os.MkdirAll(path, 0755); err != nil && !os.IsExist(err) {
		return err
	}

	if err := ioutil.WriteFile(filepath.Join(path, "cgroup.procs"), []byte(strconv.Itoa(pid)), 0700); err != nil {
		return err
	}

	if !c.AllowAllDevices {
		if err := writeFile(path, "devices.deny", "a"); err != nil {
			return err
		}

		for _, dev := range c.AllowedDevices {
			if err := writeFile(path, "devices.allow", dev.GetCgroupAllowString()); err != nil {
				return err
			}
		}
	}

	return nil
}

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

func joinMemory(c *cgroups.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 *cgroups.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "cpuset")
	if err != nil {
		return err
	}

	s := &fs.CpusetGroup{}

	return s.SetDir(path, c.CpusetCpus, c.CpusetMems, pid)
}
<commit_msg>cgroups: systemd: set DefaultDependencies=false if possible<commit_after>// +build linux

package systemd

import (
	"bytes"
	"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/godbus/dbus"
)

type systemdCgroup struct {
	cgroup *cgroups.Cgroup
}

type subsystem interface {
	GetStats(string, *cgroups.Stats) error
}

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
				}
			}
		}

		// Assume StartTransientUnit on a scope allows DefaultDependencies
		hasTransientDefaultDependencies = true
		ddf := newProp("DefaultDependencies", false)
		if _, err := theConn.StartTransientUnit("docker-systemd-test-default-dependencies.scope", "replace", ddf); err != nil {
			if dbusError, ok := err.(dbus.Error); ok {
				if dbusError.Name == "org.freedesktop.DBus.Error.PropertyReadOnly" {
					hasTransientDefaultDependencies = false
				}
			}
		}
	}
	return hasStartTransientUnit
}

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

func Apply(c *cgroups.Cgroup, pid int) (map[string]string, error) {
	var (
		unitName   = getUnitName(c)
		slice      = "system.slice"
		properties []systemd.Property
		res        = &systemdCgroup{}
	)

	res.cgroup = c

	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 nil, err
	}

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

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

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

	paths := make(map[string]string)
	for _, sysname := range []string{
		"devices",
		"memory",
		"cpu",
		"cpuset",
		"cpuacct",
		"blkio",
		"perf_event",
		"freezer",
	} {
		subsystemPath, err := getSubsystemPath(res.cgroup, sysname)
		if err != nil {
			// Don't fail if a cgroup hierarchy was not found, just skip this subsystem
			if cgroups.IsNotFound(err) {
				continue
			}
			return nil, err
		}
		paths[sysname] = subsystemPath
	}
	return paths, nil
}

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

func joinFreezer(c *cgroups.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "freezer")
	if err != nil {
		return err
	}

	if err := os.MkdirAll(path, 0755); err != nil && !os.IsExist(err) {
		return err
	}

	return ioutil.WriteFile(filepath.Join(path, "cgroup.procs"), []byte(strconv.Itoa(pid)), 0700)
}

func getSubsystemPath(c *cgroups.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 Freeze(c *cgroups.Cgroup, state cgroups.FreezerState) error {
	path, err := getSubsystemPath(c, "freezer")
	if err != nil {
		return err
	}

	if err := ioutil.WriteFile(filepath.Join(path, "freezer.state"), []byte(state), 0); err != nil {
		return err
	}
	for {
		state_, err := ioutil.ReadFile(filepath.Join(path, "freezer.state"))
		if err != nil {
			return err
		}
		if string(state) == string(bytes.TrimSpace(state_)) {
			break
		}
		time.Sleep(1 * time.Millisecond)
	}
	return nil
}

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

	return cgroups.ReadProcsFile(path)
}

func getUnitName(c *cgroups.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 *cgroups.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "devices")
	if err != nil {
		return err
	}

	if err := os.MkdirAll(path, 0755); err != nil && !os.IsExist(err) {
		return err
	}

	if err := ioutil.WriteFile(filepath.Join(path, "cgroup.procs"), []byte(strconv.Itoa(pid)), 0700); err != nil {
		return err
	}

	if !c.AllowAllDevices {
		if err := writeFile(path, "devices.deny", "a"); err != nil {
			return err
		}

		for _, dev := range c.AllowedDevices {
			if err := writeFile(path, "devices.allow", dev.GetCgroupAllowString()); err != nil {
				return err
			}
		}
	}

	return nil
}

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

func joinMemory(c *cgroups.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 *cgroups.Cgroup, pid int) error {
	path, err := getSubsystemPath(c, "cpuset")
	if err != nil {
		return err
	}

	s := &fs.CpusetGroup{}

	return s.SetDir(path, c.CpusetCpus, c.CpusetMems, pid)
}
<|endoftext|>


<commit_before>package tempconsul

import (
	"github.com/armon/consul-kv"
	"testing"
)

func TestServer(t *testing.T) {
	server := Server{}
	err := server.Start()
	defer server.Term()
	if err != nil {
		t.Fatal(err)
	}

	client, _ := consulkv.NewClient(consulkv.DefaultConfig())
	err = client.Put("test", []byte("cool"), 0)
	if err != nil {
		t.Fatal(err)
	}
	_, pair, err := client.Get("test")
	if err != nil {
		t.Fatal(err)
	}
	if string(pair.Value) != "cool" {
		t.Errorf("got: %#v", string(pair.Value))
	}
}
<commit_msg>Use new consul/api package.<commit_after>package tempconsul

import (
	"github.com/hashicorp/consul/api"
	"testing"
)

func TestServer(t *testing.T) {
	server := Server{}
	err := server.Start()
	defer server.Term()
	if err != nil {
		t.Fatal(err)
	}

	client, _ := api.NewClient(api.DefaultConfig())
	kv := client.KV()

	putpair := &api.KVPair{Key: "test", Value: []byte("cool")}
	_, err = kv.Put(putpair, nil)
	if err != nil {
		t.Fatal(err)
	}
	getpair, _, err := kv.Get("test", nil)
	if err != nil {
		t.Fatal(err)
	}
	if string(getpair.Value) != "cool" {
		t.Errorf("got: %#v", string(getpair.Value))
	}
}
<|endoftext|>


<commit_before>package myirmaserver

import (
	"sync"
	"time"

	"github.com/privacybydesign/irmago/internal/common"
	"github.com/privacybydesign/irmago/server"
)

type Sessiondata struct {
	sync.Mutex

	token  string
	userID *int64

	pendingError        *server.Error
	pendingErrorMessage string

	expiry time.Time
}

type SessionStore interface {
	create() *Sessiondata
	get(token string) *Sessiondata
	flush()
}

type MemorySessionStore struct {
	sync.Mutex

	data            map[string]*Sessiondata
	sessionLifetime time.Duration
}

func NewMemorySessionStore(sessionLifetime time.Duration) SessionStore {
	return &MemorySessionStore{
		sessionLifetime: sessionLifetime,
		data:            map[string]*Sessiondata{},
	}
}

func (s *MemorySessionStore) create() *Sessiondata {
	s.Lock()
	defer s.Unlock()
	token := common.NewSessionToken()
	s.data[token] = &Sessiondata{
		token:  token,
		expiry: time.Now().Add(s.sessionLifetime),
	}
	return s.data[token]
}

func (s *MemorySessionStore) get(token string) *Sessiondata {
	s.Lock()
	defer s.Unlock()
	return s.data[token]
}

func (s *MemorySessionStore) flush() {
	s.Lock()
	defer s.Unlock()
	for k, v := range s.data {
		if time.Now().After(v.expiry) {
			delete(s.data, k)
		}
	}
}
<commit_msg>feat: prevent repetative time.Now() computation in loop<commit_after>package myirmaserver

import (
	"sync"
	"time"

	"github.com/privacybydesign/irmago/internal/common"
	"github.com/privacybydesign/irmago/server"
)

type Sessiondata struct {
	sync.Mutex

	token  string
	userID *int64

	pendingError        *server.Error
	pendingErrorMessage string

	expiry time.Time
}

type SessionStore interface {
	create() *Sessiondata
	get(token string) *Sessiondata
	flush()
}

type MemorySessionStore struct {
	sync.Mutex

	data            map[string]*Sessiondata
	sessionLifetime time.Duration
}

func NewMemorySessionStore(sessionLifetime time.Duration) SessionStore {
	return &MemorySessionStore{
		sessionLifetime: sessionLifetime,
		data:            map[string]*Sessiondata{},
	}
}

func (s *MemorySessionStore) create() *Sessiondata {
	s.Lock()
	defer s.Unlock()
	token := common.NewSessionToken()
	s.data[token] = &Sessiondata{
		token:  token,
		expiry: time.Now().Add(s.sessionLifetime),
	}
	return s.data[token]
}

func (s *MemorySessionStore) get(token string) *Sessiondata {
	s.Lock()
	defer s.Unlock()
	return s.data[token]
}

func (s *MemorySessionStore) flush() {
	now := time.Now()
	s.Lock()
	defer s.Unlock()
	for k, v := range s.data {
		if now.After(v.expiry) {
			delete(s.data, k)
		}
	}
}
<|endoftext|>


<commit_before>package terraform

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
)

const appName string = "concourse-up"

// IClient is an interface for the terraform Client
type IClient interface {
	Output() (*Metadata, error)
	Apply() error
	Destroy() error
	Cleanup() error
}

// Client wraps common terraform commands
type Client struct {
	configDir string
	stdout    io.Writer
	stderr    io.Writer
}

// ClientFactory is a function that builds a client interface
type ClientFactory func(config []byte, stdout, stderr io.Writer) (IClient, error)

// NewClient is a concrete implementation of ClientFactory
func NewClient(config []byte, stdout, stderr io.Writer) (IClient, error) {
	if err := checkTerraformOnPath(stderr, stderr); err != nil {
		return nil, err
	}

	configDir, err := initConfig(config, stderr, stderr)
	if err != nil {
		return nil, err
	}

	return &Client{
		configDir: configDir,
		stdout:    stdout,
		stderr:    stderr,
	}, nil
}

// Cleanup cleans up the temporary directory used by terraform
func (client *Client) Cleanup() error {
	return os.RemoveAll(client.configDir)
}

// Output fetches the terraform output/metadata
func (client *Client) Output() (*Metadata, error) {
	stdoutBuffer := bytes.NewBuffer(nil)
	if err := terraform([]string{
		"output",
		"-json",
	}, client.configDir, stdoutBuffer, client.stderr); err != nil {
		return nil, err
	}

	metadata := Metadata{}
	if err := json.NewDecoder(stdoutBuffer).Decode(&metadata); err != nil {
		return nil, err
	}

	return &metadata, nil
}

// Apply takes a terraform config and applies it
func (client *Client) Apply() error {
	return terraform([]string{
		"apply",
		"-input=false",
	}, client.configDir, client.stdout, client.stderr)
}

// Destroy destroys the given terraform config
func (client *Client) Destroy() error {
	return terraform([]string{
		"destroy",
		"-force",
	}, client.configDir, client.stdout, client.stderr)
}

func checkTerraformOnPath(stdout, stderr io.Writer) error {
	if err := terraform([]string{"version"}, "", stdout, stderr); err != nil {
		return fmt.Errorf("Error running `terraform version`, is terraform in your PATH?\n%s\n\nDownload terraform here: https://www.terraform.io/downloads.html\n", err.Error())
	}
	return nil
}

func terraform(args []string, dir string, stdout, stderr io.Writer) error {
	cmd := exec.Command("terraform", args...)
	cmd.Dir = dir
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	return cmd.Run()
}

func initConfig(config []byte, stdout, stderr io.Writer) (string, error) {
	// write out config
	tmpDir, err := ioutil.TempDir("", appName)
	if err != nil {
		return "", err
	}

	configPath := filepath.Join(tmpDir, "main.tf")
	if err := ioutil.WriteFile(configPath, config, 0777); err != nil {
		return "", err
	}

	if err := terraform([]string{
		"init",
	}, tmpDir, stdout, stderr); err != nil {
		return "", err
	}

	return tmpDir, nil
}
<commit_msg>fix linting error<commit_after>package terraform

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
)

const appName string = "concourse-up"

// IClient is an interface for the terraform Client
type IClient interface {
	Output() (*Metadata, error)
	Apply() error
	Destroy() error
	Cleanup() error
}

// Client wraps common terraform commands
type Client struct {
	configDir string
	stdout    io.Writer
	stderr    io.Writer
}

// ClientFactory is a function that builds a client interface
type ClientFactory func(config []byte, stdout, stderr io.Writer) (IClient, error)

// NewClient is a concrete implementation of ClientFactory
func NewClient(config []byte, stdout, stderr io.Writer) (IClient, error) {
	if err := checkTerraformOnPath(stderr, stderr); err != nil {
		return nil, err
	}

	configDir, err := initConfig(config, stderr, stderr)
	if err != nil {
		return nil, err
	}

	return &Client{
		configDir: configDir,
		stdout:    stdout,
		stderr:    stderr,
	}, nil
}

// Cleanup cleans up the temporary directory used by terraform
func (client *Client) Cleanup() error {
	return os.RemoveAll(client.configDir)
}

// Output fetches the terraform output/metadata
func (client *Client) Output() (*Metadata, error) {
	stdoutBuffer := bytes.NewBuffer(nil)
	if err := terraform([]string{
		"output",
		"-json",
	}, client.configDir, stdoutBuffer, client.stderr); err != nil {
		return nil, err
	}

	metadata := Metadata{}
	if err := json.NewDecoder(stdoutBuffer).Decode(&metadata); err != nil {
		return nil, err
	}

	return &metadata, nil
}

// Apply takes a terraform config and applies it
func (client *Client) Apply() error {
	return terraform([]string{
		"apply",
		"-input=false",
	}, client.configDir, client.stdout, client.stderr)
}

// Destroy destroys the given terraform config
func (client *Client) Destroy() error {
	return terraform([]string{
		"destroy",
		"-force",
	}, client.configDir, client.stdout, client.stderr)
}

func checkTerraformOnPath(stdout, stderr io.Writer) error {
	if err := terraform([]string{"version"}, "", stdout, stderr); err != nil {
		return fmt.Errorf("error running `terraform version`, is terraform in your PATH?\n%s\n\nDownload terraform here: https://www.terraform.io/downloads.html\n", err.Error())
	}
	return nil
}

func terraform(args []string, dir string, stdout, stderr io.Writer) error {
	cmd := exec.Command("terraform", args...)
	cmd.Dir = dir
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	return cmd.Run()
}

func initConfig(config []byte, stdout, stderr io.Writer) (string, error) {
	// write out config
	tmpDir, err := ioutil.TempDir("", appName)
	if err != nil {
		return "", err
	}

	configPath := filepath.Join(tmpDir, "main.tf")
	if err := ioutil.WriteFile(configPath, config, 0777); err != nil {
		return "", err
	}

	if err := terraform([]string{
		"init",
	}, tmpDir, stdout, stderr); err != nil {
		return "", err
	}

	return tmpDir, nil
}
<|endoftext|>