<commit_before>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"html/template"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/zaunerc/cntrinfod/consul"
	"github.com/zaunerc/cntrinfod/docker"
	"github.com/zaunerc/cntrinfod/system"

	auth "github.com/abbot/go-http-auth"
	"github.com/shurcooL/github_flavored_markdown"
	"github.com/shurcooL/github_flavored_markdown/gfmstyle"
	"github.com/urfave/cli"
)

type Page struct {
	Path     string
	Markdown []byte
	Html     []byte
}

type Log struct {
	Path        string
	FileName    string
	Lines       string
	RecentLines string
}

func convertMdToHtml(readme []byte) (*Page, error) {

	location := locateReadme()
	enrichedReadme := enrichMd(getReadmeAsMarkdown(location))

	readmeAsHtml := github_flavored_markdown.Markdown(enrichedReadme)
	return &Page{Path: location, Markdown: enrichedReadme, Html: readmeAsHtml}, nil
}

/**
 * Adds various informations to the README.md of the container.
 */
func enrichMd(readme []byte) []byte {

	appendixTemplate, error := template.ParseFiles(path.Join(staticDataDir, "appendix_template.md"))
	var appendixBuffer bytes.Buffer

	vars := map[string]interface{}{
		"ContainerHostname": system.FetchContainerHostname(),
		"HostHostname":      docker.FetchHostHostname(),
		"TcpSocketInfo":     system.FetchTcp46SocketInfo(),
		"UdpSocketInfo":     system.FetchUdp46SocketInfo(),
		"ProcessInfo":       system.FetchProcessInfo(),
		"ProcessTree":       system.FetchProcessTree(),
	}

	error = appendixTemplate.Execute(&appendixBuffer, vars)

	if error != nil {
		fmt.Printf("Error while processing template: >%s<.", error)
		return readme
	}

	// Make sure that there is at least one newline before our heading
	readme = append(readme, "\n"...)
	return append(readme, appendixBuffer.Bytes()...)
}

// Use getReadmeAsMarkdown
// `/var/log/AdminServer` - sshd log file
func parseLogFilePathsFromMd(pathToFile string) (*[]string, error) {

	var foundPath []string

	pageAsMarkdown, error := os.Open(pathToFile)

	if error != nil {
		return nil, error
	}

	matcher := regexp.MustCompile("`(.*)`.*-.*log.*")
	scanner := bufio.NewScanner(pageAsMarkdown)
	for scanner.Scan() {
		line := scanner.Text()
		match := matcher.FindStringSubmatch(line)
		if len(match) > 0 {
			foundPath = append(foundPath, match[1])
			fmt.Printf("Found path >%s< in file >%s<.\n", match[1], pathToFile)
		}
	}

	fmt.Printf("Found %d paths in file >%s<\n", len(foundPath), pathToFile)
	return &foundPath, nil
}

func readLogFiles(pathToLogFiles *[]string) (*[]Log, error) {

	var logs []Log

	for _, pathToLogFile := range *pathToLogFiles {
		file, err := os.Open(pathToLogFile)

		if err != nil {
			fmt.Printf("Error while reading log file >%s<: >%s<\n", pathToLogFile, err)
		} else {
			fmt.Printf("Successfully read log file >%s<.\n", pathToLogFile)
		}

		defer file.Close()

		var lines []string
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			lines = append(lines, scanner.Text())
		}

		threshhold := 250
		var recentLines []string
		if len(lines) > threshhold {
			recentLines = lines[len(lines)-threshhold:]
		} else {
			recentLines = lines
		}

		log := Log{Path: pathToLogFile, FileName: filepath.Base(pathToLogFile), Lines: strings.Join(lines, "\n"), RecentLines: strings.Join(recentLines, "\n")}
		logs = append(logs, log)
	}

	return &logs, nil
}

func init() {
	// By default logger is set to write to stderr device.
	//log.SetOutput(os.Stdout)
}

var staticDataDir string

func main() {

	var httpPort int
	var htpasswd string
	var consulUrl string

	app := cli.NewApp()

	app.Flags = []cli.Flag{
		cli.IntFlag{
			Name:        "httpPort, p",
			Value:       2020,
			Usage:       "Listen on port `PORT` for HTTP connections",
			Destination: &httpPort,
		},
		cli.StringFlag{
			Name: "htpasswd, s",
			//Value:       "/usr/local/etc/cntrinfod/htpasswd",
			Usage:       "Use htpasswd `FILE`. Enables user authentication.",
			Destination: &htpasswd,
		},
		cli.StringFlag{
			Name:        "consulUrl, c",
			Usage:       "Register container with consul at `URL`. Enables consul registration. E.g. \"registry.mynet.root.local:8500\".",
			Destination: &consulUrl,
		},
		cli.StringFlag{
			Name:        "staticDataDir, d",
			Value:       "/usr/local/share/cntrinfod/",
			Usage:       "Directory containing static web site files.",
			Destination: &staticDataDir,
		},
	}

	app.Email = "christoph.zauner@NLLK.net"
	app.Author = "Christoph Zauner"
	app.Version = "0.2.4"
	// cntrinfod, cntinfod
	app.Usage = "Container Info Daemon: HTTP daemon which exposes and augments the containers REAMDE.md"

	app.Action = func(c *cli.Context) error {

		if consulUrl != "" {
			consul.ScheduleRegistration(consulUrl, httpPort)
		}

		fmt.Printf("Serving static files from >%s<\n", getAbsolutePath(staticDataDir))
		fmt.Printf("Starting HTTP daemon on port %d...\n", httpPort)

		http.HandleFunc("/", protect(handler, htpasswd))
		http.HandleFunc("/log", protect(logHandler, htpasswd))
		http.HandleFunc("/hostinfo", protect(hostInfoHandler, htpasswd))
		http.HandleFunc("/markdown", protect(markdownHandler, htpasswd))

		// Serve the "/assets/gfm.css" file.
		http.Handle("/assets/", http.StripPrefix("/assets", http.FileServer(gfmstyle.Assets)))

		http.ListenAndServe(":"+strconv.Itoa(httpPort), nil)

		return nil
	}

	app.Run(os.Args)
}

func protect(handlerFunc http.HandlerFunc, htpasswdPath string) http.HandlerFunc {

	if htpasswdPath != "" {
		// read from htpasswd file
		htpasswd := auth.HtpasswdFileProvider(htpasswdPath)
		authenticator := auth.NewBasicAuthenticator("cntrinfod htpasswd realm", htpasswd)
		return auth.JustCheck(authenticator, handlerFunc)
	} else {
		return func(w http.ResponseWriter, r *http.Request) {
			handlerFunc(w, r)
		}
	}
}

func locateReadme() string {

	var possibleLocations [2]string = [2]string{"/README.md",
		"testdata/README.md"}

	var fd *os.File
	defer fd.Close()
	var err error

	var location string
	for _, possibleLocation := range possibleLocations {
		fd, err = os.Open(possibleLocation)
		if err == nil {
			fmt.Printf("Found README.md file: >%s<\n", possibleLocation)
			location = possibleLocation
			break
		} else {
			fmt.Printf("Error when looking for README.md at >%s<: >%s<\n", possibleLocation, err)
		}
	}

	if err != nil {
		fmt.Printf("Could not find README.md file.\n")
		os.Exit(-1)
	}

	return location
}

func getReadmeAsMarkdown(path string) []byte {

	var fd *os.File
	defer fd.Close()
	var err error

	fd, err = os.Open(path)
	if err != nil {
		fmt.Printf("Error while trying to open REAMDE.md file: %s\n", err)
		os.Exit(-1)
	}

	var readme []byte
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		readme = append(readme, scanner.Bytes()...)
		readme = append(readme, "\n"...)
	}

	return readme
}

func handler(w http.ResponseWriter, r *http.Request) {
	page, _ := convertMdToHtml(getReadmeAsMarkdown(locateReadme()))

	io.WriteString(w, `<meta charset="utf-8"><link href="/assets/gfm.css" media="all" rel="stylesheet" type="text/css" /><link href="//cdnjs.cloudflare.com/ajax/libs/octicons/2.1.2/octicons.css" media="all" rel="stylesheet" type="text/css" />
`)
	w.Write(page.Html)
	io.WriteString(w, `</article></body></html>`)
}

func logHandler(w http.ResponseWriter, r *http.Request) {

	pathToLogFiles, _ := parseLogFilePathsFromMd(locateReadme())
	logs, _ := readLogFiles(pathToLogFiles)

	t, error := template.ParseFiles(path.Join(staticDataDir, "log.html"))

	currentDateAndTime := time.Now().Format("2006-01-02 15:04:05")

	vars := map[string]interface{}{
		"Logs":               logs,
		"CurrentDateAndTime": currentDateAndTime,
	}

	error = t.Execute(w, vars)

	if error != nil {
		fmt.Printf("Error while processing template: >%s<.", error)
	}
}

func hostInfoHandler(w http.ResponseWriter, r *http.Request) {
	t, error := template.ParseFiles(path.Join(staticDataDir, "hostinfo.html"))

	hostinfo := docker.FetchHostInfo()
	error = t.Execute(w, hostinfo)

	if error != nil {
		fmt.Printf("Error while processing template: >%s<.", error)
	}
}

func markdownHandler(w http.ResponseWriter, r *http.Request) {

	readmeAsMarkdown := getReadmeAsMarkdown(locateReadme())
	enrichedReadme := enrichMd(readmeAsMarkdown)

	t, error := template.ParseFiles(path.Join(staticDataDir, "markdown.html"))
	error = t.Execute(w, string(enrichedReadme))
	if error != nil {
		fmt.Printf("Error while processing template: >%s<.", error)
	}
}

func getAbsolutePath(name string) string {
	if path.IsAbs(name) {
		return name
	}
	wd, err := os.Getwd()

	if err != nil {
		panic(fmt.Sprintf("Error while trying to get current working dir: %s", err))
	}

	return path.Join(wd, name)
}
<commit_msg>Starting development iteration on 0.2.5-SNAPSHOT.<commit_after>package main

import (
	"bufio"
	"bytes"
	"fmt"
	"html/template"
	"io"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/zaunerc/cntrinfod/consul"
	"github.com/zaunerc/cntrinfod/docker"
	"github.com/zaunerc/cntrinfod/system"

	auth "github.com/abbot/go-http-auth"
	"github.com/shurcooL/github_flavored_markdown"
	"github.com/shurcooL/github_flavored_markdown/gfmstyle"
	"github.com/urfave/cli"
)

type Page struct {
	Path     string
	Markdown []byte
	Html     []byte
}

type Log struct {
	Path        string
	FileName    string
	Lines       string
	RecentLines string
}

func convertMdToHtml(readme []byte) (*Page, error) {

	location := locateReadme()
	enrichedReadme := enrichMd(getReadmeAsMarkdown(location))

	readmeAsHtml := github_flavored_markdown.Markdown(enrichedReadme)
	return &Page{Path: location, Markdown: enrichedReadme, Html: readmeAsHtml}, nil
}

/**
 * Adds various informations to the README.md of the container.
 */
func enrichMd(readme []byte) []byte {

	appendixTemplate, error := template.ParseFiles(path.Join(staticDataDir, "appendix_template.md"))
	var appendixBuffer bytes.Buffer

	vars := map[string]interface{}{
		"ContainerHostname": system.FetchContainerHostname(),
		"HostHostname":      docker.FetchHostHostname(),
		"TcpSocketInfo":     system.FetchTcp46SocketInfo(),
		"UdpSocketInfo":     system.FetchUdp46SocketInfo(),
		"ProcessInfo":       system.FetchProcessInfo(),
		"ProcessTree":       system.FetchProcessTree(),
	}

	error = appendixTemplate.Execute(&appendixBuffer, vars)

	if error != nil {
		fmt.Printf("Error while processing template: >%s<.", error)
		return readme
	}

	// Make sure that there is at least one newline before our heading
	readme = append(readme, "\n"...)
	return append(readme, appendixBuffer.Bytes()...)
}

// Use getReadmeAsMarkdown
// `/var/log/AdminServer` - sshd log file
func parseLogFilePathsFromMd(pathToFile string) (*[]string, error) {

	var foundPath []string

	pageAsMarkdown, error := os.Open(pathToFile)

	if error != nil {
		return nil, error
	}

	matcher := regexp.MustCompile("`(.*)`.*-.*log.*")
	scanner := bufio.NewScanner(pageAsMarkdown)
	for scanner.Scan() {
		line := scanner.Text()
		match := matcher.FindStringSubmatch(line)
		if len(match) > 0 {
			foundPath = append(foundPath, match[1])
			fmt.Printf("Found path >%s< in file >%s<.\n", match[1], pathToFile)
		}
	}

	fmt.Printf("Found %d paths in file >%s<\n", len(foundPath), pathToFile)
	return &foundPath, nil
}

func readLogFiles(pathToLogFiles *[]string) (*[]Log, error) {

	var logs []Log

	for _, pathToLogFile := range *pathToLogFiles {
		file, err := os.Open(pathToLogFile)

		if err != nil {
			fmt.Printf("Error while reading log file >%s<: >%s<\n", pathToLogFile, err)
		} else {
			fmt.Printf("Successfully read log file >%s<.\n", pathToLogFile)
		}

		defer file.Close()

		var lines []string
		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			lines = append(lines, scanner.Text())
		}

		threshhold := 250
		var recentLines []string
		if len(lines) > threshhold {
			recentLines = lines[len(lines)-threshhold:]
		} else {
			recentLines = lines
		}

		log := Log{Path: pathToLogFile, FileName: filepath.Base(pathToLogFile), Lines: strings.Join(lines, "\n"), RecentLines: strings.Join(recentLines, "\n")}
		logs = append(logs, log)
	}

	return &logs, nil
}

func init() {
	// By default logger is set to write to stderr device.
	//log.SetOutput(os.Stdout)
}

var staticDataDir string

func main() {

	var httpPort int
	var htpasswd string
	var consulUrl string

	app := cli.NewApp()

	app.Flags = []cli.Flag{
		cli.IntFlag{
			Name:        "httpPort, p",
			Value:       2020,
			Usage:       "Listen on port `PORT` for HTTP connections",
			Destination: &httpPort,
		},
		cli.StringFlag{
			Name: "htpasswd, s",
			//Value:       "/usr/local/etc/cntrinfod/htpasswd",
			Usage:       "Use htpasswd `FILE`. Enables user authentication.",
			Destination: &htpasswd,
		},
		cli.StringFlag{
			Name:        "consulUrl, c",
			Usage:       "Register container with consul at `URL`. Enables consul registration. E.g. \"registry.mynet.root.local:8500\".",
			Destination: &consulUrl,
		},
		cli.StringFlag{
			Name:        "staticDataDir, d",
			Value:       "/usr/local/share/cntrinfod/",
			Usage:       "Directory containing static web site files.",
			Destination: &staticDataDir,
		},
	}

	app.Email = "christoph.zauner@NLLK.net"
	app.Author = "Christoph Zauner"
	app.Version = "0.2.5-SNAPSHOT"
	// cntrinfod, cntinfod
	app.Usage = "Container Info Daemon: HTTP daemon which exposes and augments the containers REAMDE.md"

	app.Action = func(c *cli.Context) error {

		if consulUrl != "" {
			consul.ScheduleRegistration(consulUrl, httpPort)
		}

		fmt.Printf("Serving static files from >%s<\n", getAbsolutePath(staticDataDir))
		fmt.Printf("Starting HTTP daemon on port %d...\n", httpPort)

		http.HandleFunc("/", protect(handler, htpasswd))
		http.HandleFunc("/log", protect(logHandler, htpasswd))
		http.HandleFunc("/hostinfo", protect(hostInfoHandler, htpasswd))
		http.HandleFunc("/markdown", protect(markdownHandler, htpasswd))

		// Serve the "/assets/gfm.css" file.
		http.Handle("/assets/", http.StripPrefix("/assets", http.FileServer(gfmstyle.Assets)))

		http.ListenAndServe(":"+strconv.Itoa(httpPort), nil)

		return nil
	}

	app.Run(os.Args)
}

func protect(handlerFunc http.HandlerFunc, htpasswdPath string) http.HandlerFunc {

	if htpasswdPath != "" {
		// read from htpasswd file
		htpasswd := auth.HtpasswdFileProvider(htpasswdPath)
		authenticator := auth.NewBasicAuthenticator("cntrinfod htpasswd realm", htpasswd)
		return auth.JustCheck(authenticator, handlerFunc)
	} else {
		return func(w http.ResponseWriter, r *http.Request) {
			handlerFunc(w, r)
		}
	}
}

func locateReadme() string {

	var possibleLocations [2]string = [2]string{"/README.md",
		"testdata/README.md"}

	var fd *os.File
	defer fd.Close()
	var err error

	var location string
	for _, possibleLocation := range possibleLocations {
		fd, err = os.Open(possibleLocation)
		if err == nil {
			fmt.Printf("Found README.md file: >%s<\n", possibleLocation)
			location = possibleLocation
			break
		} else {
			fmt.Printf("Error when looking for README.md at >%s<: >%s<\n", possibleLocation, err)
		}
	}

	if err != nil {
		fmt.Printf("Could not find README.md file.\n")
		os.Exit(-1)
	}

	return location
}

func getReadmeAsMarkdown(path string) []byte {

	var fd *os.File
	defer fd.Close()
	var err error

	fd, err = os.Open(path)
	if err != nil {
		fmt.Printf("Error while trying to open REAMDE.md file: %s\n", err)
		os.Exit(-1)
	}

	var readme []byte
	scanner := bufio.NewScanner(fd)
	for scanner.Scan() {
		readme = append(readme, scanner.Bytes()...)
		readme = append(readme, "\n"...)
	}

	return readme
}

func handler(w http.ResponseWriter, r *http.Request) {
	page, _ := convertMdToHtml(getReadmeAsMarkdown(locateReadme()))

	io.WriteString(w, `<meta charset="utf-8"><link href="/assets/gfm.css" media="all" rel="stylesheet" type="text/css" /><link href="//cdnjs.cloudflare.com/ajax/libs/octicons/2.1.2/octicons.css" media="all" rel="stylesheet" type="text/css" />
`)
	w.Write(page.Html)
	io.WriteString(w, `</article></body></html>`)
}

func logHandler(w http.ResponseWriter, r *http.Request) {

	pathToLogFiles, _ := parseLogFilePathsFromMd(locateReadme())
	logs, _ := readLogFiles(pathToLogFiles)

	t, error := template.ParseFiles(path.Join(staticDataDir, "log.html"))

	currentDateAndTime := time.Now().Format("2006-01-02 15:04:05")

	vars := map[string]interface{}{
		"Logs":               logs,
		"CurrentDateAndTime": currentDateAndTime,
	}

	error = t.Execute(w, vars)

	if error != nil {
		fmt.Printf("Error while processing template: >%s<.", error)
	}
}

func hostInfoHandler(w http.ResponseWriter, r *http.Request) {
	t, error := template.ParseFiles(path.Join(staticDataDir, "hostinfo.html"))

	hostinfo := docker.FetchHostInfo()
	error = t.Execute(w, hostinfo)

	if error != nil {
		fmt.Printf("Error while processing template: >%s<.", error)
	}
}

func markdownHandler(w http.ResponseWriter, r *http.Request) {

	readmeAsMarkdown := getReadmeAsMarkdown(locateReadme())
	enrichedReadme := enrichMd(readmeAsMarkdown)

	t, error := template.ParseFiles(path.Join(staticDataDir, "markdown.html"))
	error = t.Execute(w, string(enrichedReadme))
	if error != nil {
		fmt.Printf("Error while processing template: >%s<.", error)
	}
}

func getAbsolutePath(name string) string {
	if path.IsAbs(name) {
		return name
	}
	wd, err := os.Getwd()

	if err != nil {
		panic(fmt.Sprintf("Error while trying to get current working dir: %s", err))
	}

	return path.Join(wd, name)
}
<|endoftext|>


<commit_before>package code

import (
	"net/http"
)

type (
	Code         int
	TextHTTPCode struct {
		Text     string
		HTTPCode int
	}
	CodeMap map[Code]TextHTTPCode
)

const (
	BalanceNoEnough  Code = -10 //余额不足
	RequestTimeout   Code = -9  //提交超时
	AbnormalResponse Code = -8  //响应异常
	OperationTimeout Code = -7  //操作超时
	Unsupported      Code = -6  //不支持的操作
	RepeatOperation  Code = -5  //重复操作
	DataNotFound     Code = -4  //数据未找到
	UserNotFound     Code = -3  //用户未找到
	NonPrivileged    Code = -2  //无权限
	Unauthenticated  Code = -1  //未登录
	Failure          Code = 0   //操作失败
	Success          Code = 1   //操作成功
)

// CodeDict 状态码字典
var CodeDict = CodeMap{
	BalanceNoEnough:  {"BalanceNoEnough", http.StatusOK},
	RequestTimeout:   {"RequestTimeout", http.StatusOK},
	AbnormalResponse: {"AbnormalResponse", http.StatusOK},
	OperationTimeout: {"OperationTimeout", http.StatusOK},
	Unsupported:      {"Unsupported", http.StatusOK},
	RepeatOperation:  {"RepeatOperation", http.StatusOK},
	DataNotFound:     {"DataNotFound", http.StatusOK},
	UserNotFound:     {"UserNotFound", http.StatusOK},
	NonPrivileged:    {"NonPrivileged", http.StatusOK},
	Unauthenticated:  {"Unauthenticated", http.StatusOK},
	Failure:          {"Failure", http.StatusOK},
	Success:          {"Success", http.StatusOK},
}

func (c Code) String() string {
	if v, y := CodeDict[c]; y {
		return v.Text
	}
	return `Undefined`
}

// Int 返回int类型的自定义状态码
func (c Code) Int() int {
	return int(c)
}

// HTTPCode 返回HTTP状态码
func (c Code) HTTPCode() int {
	if v, y := CodeDict[c]; y {
		return v.HTTPCode
	}
	return http.StatusOK
}

func (s CodeMap) Get(code int) TextHTTPCode {
	v, _ := s[Code(code)]
	return v
}
<commit_msg>improved<commit_after>package code

import (
	"net/http"
)

type (
	Code         int
	TextHTTPCode struct {
		Text     string
		HTTPCode int
	}
	CodeMap map[Code]TextHTTPCode
)

const (
	// - 操作状态

	RequestTimeout   Code = -204  //提交超时
	AbnormalResponse Code = -203  //响应异常
	OperationTimeout Code = -202  //操作超时
	Unsupported      Code = -201  //不支持的操作
	RepeatOperation  Code = -200  //重复操作

	// - 数据状态

	BalanceNoEnough  Code = -101 //余额不足
	DataNotFound     Code = -100  //数据未找到

	// - 用户状态

	UserDisabled     Code = -4  //用户被禁用
	UserNotFound     Code = -3  //用户未找到
	NonPrivileged    Code = -2  //无权限
	Unauthenticated  Code = -1  //未登录

	// - 通用

	Failure          Code = 0   //操作失败
	Success          Code = 1   //操作成功
)

// CodeDict 状态码字典
var CodeDict = CodeMap{
	BalanceNoEnough:  {"BalanceNoEnough", http.StatusOK},
	RequestTimeout:   {"RequestTimeout", http.StatusOK},
	AbnormalResponse: {"AbnormalResponse", http.StatusOK},
	OperationTimeout: {"OperationTimeout", http.StatusOK},
	Unsupported:      {"Unsupported", http.StatusOK},
	RepeatOperation:  {"RepeatOperation", http.StatusOK},
	DataNotFound:     {"DataNotFound", http.StatusOK},
	UserNotFound:     {"UserNotFound", http.StatusOK},
	NonPrivileged:    {"NonPrivileged", http.StatusOK},
	Unauthenticated:  {"Unauthenticated", http.StatusOK},
	Failure:          {"Failure", http.StatusOK},
	Success:          {"Success", http.StatusOK},
}

func (c Code) String() string {
	if v, y := CodeDict[c]; y {
		return v.Text
	}
	return `Undefined`
}

// Int 返回int类型的自定义状态码
func (c Code) Int() int {
	return int(c)
}

// HTTPCode 返回HTTP状态码
func (c Code) HTTPCode() int {
	if v, y := CodeDict[c]; y {
		return v.HTTPCode
	}
	return http.StatusOK
}

func (s CodeMap) Get(code int) TextHTTPCode {
	v, _ := s[Code(code)]
	return v
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"context"
	"fmt"
	"log"
	"time"

	resources "github.com/Azure/azure-sdk-for-go/services/resources/mgmt/2018-05-01/resources"
	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/adal"
	"github.com/Azure/go-autorest/autorest/azure"
)

type AcsEngineAPIModel struct {
	Location   string            `json:"location,omitempty"`
	Name       string            `json:"name,omitempty"`
	Tags       map[string]string `json:"tags,omitempty"`
	APIVersion string            `json:"APIVersion"`

	Properties *Properties `json:"properties"`
}

type Properties struct {
	OrchestratorProfile     *OrchestratorProfile     `json:"orchestratorProfile,omitempty"`
	MasterProfile           *MasterProfile           `json:"masterProfile,omitempty"`
	AgentPoolProfiles       []*AgentPoolProfile      `json:"agentPoolProfiles,omitempty"`
	LinuxProfile            *LinuxProfile            `json:"linuxProfile,omitempty"`
	WindowsProfile          *WindowsProfile          `json:"windowsProfile,omitempty"`
	ServicePrincipalProfile *ServicePrincipalProfile `json:"servicePrincipalProfile,omitempty"`
	ExtensionProfiles       []map[string]string      `json:"extensionProfiles,omitempty"`
	CustomCloudProfile      *CustomCloudProfile      `json:"customCloudProfile,omitempty"`
}

type ServicePrincipalProfile struct {
	ClientID string `json:"clientId,omitempty"`
	Secret   string `json:"secret,omitempty"`
}

type LinuxProfile struct {
	AdminUsername string `json:"adminUsername"`
	SSHKeys       *SSH   `json:"ssh"`
}

type SSH struct {
	PublicKeys []PublicKey `json:"publicKeys"`
}

type PublicKey struct {
	KeyData string `json:"keyData"`
}

type WindowsProfile struct {
	AdminUsername         string `json:"adminUsername,omitempty"`
	AdminPassword         string `json:"adminPassword,omitempty"`
	ImageVersion          string `json:"imageVersion,omitempty"`
	WindowsImageSourceURL string `json:"WindowsImageSourceUrl"`
	WindowsPublisher      string `json:"WindowsPublisher"`
	WindowsOffer          string `json:"WindowsOffer"`
	WindowsSku            string `json:"WindowsSku"`
	WindowsDockerVersion  string `json:"windowsDockerVersion"`
	SSHEnabled            bool   `json:"sshEnabled,omitempty"`
}

// KubernetesContainerSpec defines configuration for a container spec
type KubernetesContainerSpec struct {
	Name           string `json:"name,omitempty"`
	Image          string `json:"image,omitempty"`
	CPURequests    string `json:"cpuRequests,omitempty"`
	MemoryRequests string `json:"memoryRequests,omitempty"`
	CPULimits      string `json:"cpuLimits,omitempty"`
	MemoryLimits   string `json:"memoryLimits,omitempty"`
}

// KubernetesAddon defines a list of addons w/ configuration to include with the cluster deployment
type KubernetesAddon struct {
	Name       string                    `json:"name,omitempty"`
	Enabled    *bool                     `json:"enabled,omitempty"`
	Containers []KubernetesContainerSpec `json:"containers,omitempty"`
	Config     map[string]string         `json:"config,omitempty"`
	Data       string                    `json:"data,omitempty"`
}

type KubernetesConfig struct {
	CustomWindowsPackageURL      string            `json:"customWindowsPackageURL,omitempty"`
	CustomHyperkubeImage         string            `json:"customHyperkubeImage,omitempty"`
	CustomCcmImage               string            `json:"customCcmImage,omitempty"` // Image for cloud-controller-manager
	UseCloudControllerManager    *bool             `json:"useCloudControllerManager,omitempty"`
	NetworkPlugin                string            `json:"networkPlugin,omitempty"`
	PrivateAzureRegistryServer   string            `json:"privateAzureRegistryServer,omitempty"`
	AzureCNIURLLinux             string            `json:"azureCNIURLLinux,omitempty"`
	AzureCNIURLWindows           string            `json:"azureCNIURLWindows,omitempty"`
	Addons                       []KubernetesAddon `json:"addons,omitempty"`
	NetworkPolicy                string            `json:"networkPolicy,omitempty"`
	CloudProviderRateLimitQPS    float64           `json:"cloudProviderRateLimitQPS,omitempty"`
	CloudProviderRateLimitBucket int               `json:"cloudProviderRateLimitBucket,omitempty"`
	APIServerConfig              map[string]string `json:"apiServerConfig,omitempty"`
	KubernetesImageBase          string            `json:"kubernetesImageBase,omitempty"`
}
type OrchestratorProfile struct {
	OrchestratorType    string            `json:"orchestratorType"`
	OrchestratorRelease string            `json:"orchestratorRelease"`
	KubernetesConfig    *KubernetesConfig `json:"kubernetesConfig,omitempty"`
}

type MasterProfile struct {
	Count          int                 `json:"count"`
	Distro         string              `json:"distro"`
	DNSPrefix      string              `json:"dnsPrefix"`
	VMSize         string              `json:"vmSize" validate:"required"`
	IPAddressCount int                 `json:"ipAddressCount,omitempty"`
	Extensions     []map[string]string `json:"extensions,omitempty"`
	OSDiskSizeGB   int                 `json:"osDiskSizeGB,omitempty" validate:"min=0,max=1023"`
}

type AgentPoolProfile struct {
	Name                   string              `json:"name"`
	Count                  int                 `json:"count"`
	Distro                 string              `json:"distro"`
	VMSize                 string              `json:"vmSize"`
	OSType                 string              `json:"osType,omitempty"`
	AvailabilityProfile    string              `json:"availabilityProfile"`
	IPAddressCount         int                 `json:"ipAddressCount,omitempty"`
	PreProvisionExtension  map[string]string   `json:"preProvisionExtension,omitempty"`
	Extensions             []map[string]string `json:"extensions,omitempty"`
	OSDiskSizeGB           int                 `json:"osDiskSizeGB,omitempty" validate:"min=0,max=1023"`
	EnableVMSSNodePublicIP bool                `json:"enableVMSSNodePublicIP,omitempty"`
}

type AzureClient struct {
	environment       azure.Environment
	subscriptionID    string
	deploymentsClient resources.DeploymentsClient
	groupsClient      resources.GroupsClient
}

// CustomCloudProfile defines configuration for custom cloud profile( for ex: Azure Stack)
type CustomCloudProfile struct {
	PortalURL string `json:"portalURL,omitempty"`
}

// AzureStackMetadataEndpoints defines configuration for Azure Stack
type AzureStackMetadataEndpoints struct {
	GalleryEndpoint string                            `json:"galleryEndpoint,omitempty"`
	GraphEndpoint   string                            `json:"graphEndpoint,omitempty"`
	PortalEndpoint  string                            `json:"portalEndpoint,omitempty"`
	Authentication  *AzureStackMetadataAuthentication `json:"authentication,omitempty"`
}

// AzureStackMetadataAuthentication defines configuration for Azure Stack
type AzureStackMetadataAuthentication struct {
	LoginEndpoint string   `json:"loginEndpoint,omitempty"`
	Audiences     []string `json:"audiences,omitempty"`
}

func (az *AzureClient) ValidateDeployment(ctx context.Context, resourceGroupName, deploymentName string, template, params *map[string]interface{}) (valid resources.DeploymentValidateResult, err error) {
	return az.deploymentsClient.Validate(ctx,
		resourceGroupName,
		deploymentName,
		resources.Deployment{
			Properties: &resources.DeploymentProperties{
				Template:   template,
				Parameters: params,
				Mode:       resources.Incremental,
			},
		})
}

func (az *AzureClient) DeployTemplate(ctx context.Context, resourceGroupName, deploymentName string, template, parameters *map[string]interface{}) (de resources.DeploymentExtended, err error) {
	future, err := az.deploymentsClient.CreateOrUpdate(
		ctx,
		resourceGroupName,
		deploymentName,
		resources.Deployment{
			Properties: &resources.DeploymentProperties{
				Template:   template,
				Parameters: parameters,
				Mode:       resources.Incremental,
			},
		})
	if err != nil {
		return de, fmt.Errorf("cannot create deployment: %v", err)
	}

	err = future.WaitForCompletionRef(ctx, az.deploymentsClient.Client)
	if err != nil {
		return de, fmt.Errorf("cannot get the create deployment future response: %v", err)
	}

	return future.Result(az.deploymentsClient)
}

func (az *AzureClient) EnsureResourceGroup(ctx context.Context, name, location string, managedBy *string) (resourceGroup *resources.Group, err error) {
	var tags map[string]*string
	group, err := az.groupsClient.Get(ctx, name)
	if err == nil {
		tags = group.Tags
	}

	response, err := az.groupsClient.CreateOrUpdate(ctx, name, resources.Group{
		Name:      &name,
		Location:  &location,
		ManagedBy: managedBy,
		Tags:      tags,
	})
	if err != nil {
		return &response, err
	}

	return &response, nil
}

func (az *AzureClient) DeleteResourceGroup(ctx context.Context, groupName string) error {
	_, err := az.groupsClient.Get(ctx, groupName)
	if err == nil {
		future, err := az.groupsClient.Delete(ctx, groupName)
		if err != nil {
			return fmt.Errorf("cannot delete resource group %v: %v", groupName, err)
		}
		err = future.WaitForCompletionRef(ctx, az.groupsClient.Client)
		if err != nil {
			// Skip the teardown errors because of https://github.com/Azure/go-autorest/issues/357
			// TODO(feiskyer): fix the issue by upgrading go-autorest version >= v11.3.2.
			log.Printf("Warning: failed to delete resource group %q with error %v", groupName, err)
		}
	}
	return nil
}

func getOAuthConfig(env azure.Environment, subscriptionID, tenantID string) (*adal.OAuthConfig, error) {

	oauthConfig, err := adal.NewOAuthConfig(env.ActiveDirectoryEndpoint, tenantID)
	if err != nil {
		return nil, err
	}

	return oauthConfig, nil
}

func getAzureClient(env azure.Environment, subscriptionID, clientID, tenantID, clientSecret string) (*AzureClient, error) {
	oauthConfig, err := getOAuthConfig(env, subscriptionID, tenantID)
	if err != nil {
		return nil, err
	}

	armSpt, err := adal.NewServicePrincipalToken(*oauthConfig, clientID, clientSecret, env.ServiceManagementEndpoint)
	if err != nil {
		return nil, err
	}

	return getClient(env, subscriptionID, tenantID, armSpt), nil
}

func getClient(env azure.Environment, subscriptionID, tenantID string, armSpt *adal.ServicePrincipalToken) *AzureClient {
	c := &AzureClient{
		environment:    env,
		subscriptionID: subscriptionID,

		deploymentsClient: resources.NewDeploymentsClientWithBaseURI(env.ResourceManagerEndpoint, subscriptionID),
		groupsClient:      resources.NewGroupsClientWithBaseURI(env.ResourceManagerEndpoint, subscriptionID),
	}

	authorizer := autorest.NewBearerAuthorizer(armSpt)
	c.deploymentsClient.Authorizer = authorizer
	c.deploymentsClient.PollingDuration = 60 * time.Minute
	c.groupsClient.Authorizer = authorizer

	return c
}
<commit_msg>Update azure deployer for supporting dualstack<commit_after>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"context"
	"fmt"
	"log"
	"time"

	resources "github.com/Azure/azure-sdk-for-go/services/resources/mgmt/2018-05-01/resources"
	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/adal"
	"github.com/Azure/go-autorest/autorest/azure"
)

type AcsEngineAPIModel struct {
	Location   string            `json:"location,omitempty"`
	Name       string            `json:"name,omitempty"`
	Tags       map[string]string `json:"tags,omitempty"`
	APIVersion string            `json:"APIVersion"`

	Properties *Properties `json:"properties"`
}

type Properties struct {
	OrchestratorProfile     *OrchestratorProfile     `json:"orchestratorProfile,omitempty"`
	MasterProfile           *MasterProfile           `json:"masterProfile,omitempty"`
	AgentPoolProfiles       []*AgentPoolProfile      `json:"agentPoolProfiles,omitempty"`
	LinuxProfile            *LinuxProfile            `json:"linuxProfile,omitempty"`
	WindowsProfile          *WindowsProfile          `json:"windowsProfile,omitempty"`
	ServicePrincipalProfile *ServicePrincipalProfile `json:"servicePrincipalProfile,omitempty"`
	ExtensionProfiles       []map[string]string      `json:"extensionProfiles,omitempty"`
	CustomCloudProfile      *CustomCloudProfile      `json:"customCloudProfile,omitempty"`
	FeatureFlags            *FeatureFlags            `json:"featureFlags,omitempty"`
}

type ServicePrincipalProfile struct {
	ClientID string `json:"clientId,omitempty"`
	Secret   string `json:"secret,omitempty"`
}

type LinuxProfile struct {
	AdminUsername string `json:"adminUsername"`
	SSHKeys       *SSH   `json:"ssh"`
}

type SSH struct {
	PublicKeys []PublicKey `json:"publicKeys"`
}

type PublicKey struct {
	KeyData string `json:"keyData"`
}

type WindowsProfile struct {
	AdminUsername         string `json:"adminUsername,omitempty"`
	AdminPassword         string `json:"adminPassword,omitempty"`
	ImageVersion          string `json:"imageVersion,omitempty"`
	WindowsImageSourceURL string `json:"WindowsImageSourceUrl"`
	WindowsPublisher      string `json:"WindowsPublisher"`
	WindowsOffer          string `json:"WindowsOffer"`
	WindowsSku            string `json:"WindowsSku"`
	WindowsDockerVersion  string `json:"windowsDockerVersion"`
	SSHEnabled            bool   `json:"sshEnabled,omitempty"`
}

// KubernetesContainerSpec defines configuration for a container spec
type KubernetesContainerSpec struct {
	Name           string `json:"name,omitempty"`
	Image          string `json:"image,omitempty"`
	CPURequests    string `json:"cpuRequests,omitempty"`
	MemoryRequests string `json:"memoryRequests,omitempty"`
	CPULimits      string `json:"cpuLimits,omitempty"`
	MemoryLimits   string `json:"memoryLimits,omitempty"`
}

// KubernetesAddon defines a list of addons w/ configuration to include with the cluster deployment
type KubernetesAddon struct {
	Name       string                    `json:"name,omitempty"`
	Enabled    *bool                     `json:"enabled,omitempty"`
	Containers []KubernetesContainerSpec `json:"containers,omitempty"`
	Config     map[string]string         `json:"config,omitempty"`
	Data       string                    `json:"data,omitempty"`
}

type KubernetesConfig struct {
	CustomWindowsPackageURL      string            `json:"customWindowsPackageURL,omitempty"`
	CustomHyperkubeImage         string            `json:"customHyperkubeImage,omitempty"`
	CustomCcmImage               string            `json:"customCcmImage,omitempty"` // Image for cloud-controller-manager
	UseCloudControllerManager    *bool             `json:"useCloudControllerManager,omitempty"`
	NetworkPlugin                string            `json:"networkPlugin,omitempty"`
	PrivateAzureRegistryServer   string            `json:"privateAzureRegistryServer,omitempty"`
	AzureCNIURLLinux             string            `json:"azureCNIURLLinux,omitempty"`
	AzureCNIURLWindows           string            `json:"azureCNIURLWindows,omitempty"`
	Addons                       []KubernetesAddon `json:"addons,omitempty"`
	NetworkPolicy                string            `json:"networkPolicy,omitempty"`
	CloudProviderRateLimitQPS    float64           `json:"cloudProviderRateLimitQPS,omitempty"`
	CloudProviderRateLimitBucket int               `json:"cloudProviderRateLimitBucket,omitempty"`
	APIServerConfig              map[string]string `json:"apiServerConfig,omitempty"`
	KubernetesImageBase          string            `json:"kubernetesImageBase,omitempty"`
	ControllerManagerConfig      map[string]string `json:"controllerManagerConfig,omitempty"`
	KubeletConfig                map[string]string `json:"kubeletConfig,omitempty"`
}

type OrchestratorProfile struct {
	OrchestratorType    string            `json:"orchestratorType"`
	OrchestratorRelease string            `json:"orchestratorRelease"`
	KubernetesConfig    *KubernetesConfig `json:"kubernetesConfig,omitempty"`
}

type MasterProfile struct {
	Count          int                 `json:"count"`
	Distro         string              `json:"distro"`
	DNSPrefix      string              `json:"dnsPrefix"`
	VMSize         string              `json:"vmSize" validate:"required"`
	IPAddressCount int                 `json:"ipAddressCount,omitempty"`
	Extensions     []map[string]string `json:"extensions,omitempty"`
	OSDiskSizeGB   int                 `json:"osDiskSizeGB,omitempty" validate:"min=0,max=1023"`
}

type AgentPoolProfile struct {
	Name                   string              `json:"name"`
	Count                  int                 `json:"count"`
	Distro                 string              `json:"distro"`
	VMSize                 string              `json:"vmSize"`
	OSType                 string              `json:"osType,omitempty"`
	AvailabilityProfile    string              `json:"availabilityProfile"`
	IPAddressCount         int                 `json:"ipAddressCount,omitempty"`
	PreProvisionExtension  map[string]string   `json:"preProvisionExtension,omitempty"`
	Extensions             []map[string]string `json:"extensions,omitempty"`
	OSDiskSizeGB           int                 `json:"osDiskSizeGB,omitempty" validate:"min=0,max=1023"`
	EnableVMSSNodePublicIP bool                `json:"enableVMSSNodePublicIP,omitempty"`
}

type AzureClient struct {
	environment       azure.Environment
	subscriptionID    string
	deploymentsClient resources.DeploymentsClient
	groupsClient      resources.GroupsClient
}

type FeatureFlags struct {
	EnableIPv6DualStack bool `json:"enableIPv6DualStack,omitempty"`
}

// CustomCloudProfile defines configuration for custom cloud profile( for ex: Azure Stack)
type CustomCloudProfile struct {
	PortalURL string `json:"portalURL,omitempty"`
}

// AzureStackMetadataEndpoints defines configuration for Azure Stack
type AzureStackMetadataEndpoints struct {
	GalleryEndpoint string                            `json:"galleryEndpoint,omitempty"`
	GraphEndpoint   string                            `json:"graphEndpoint,omitempty"`
	PortalEndpoint  string                            `json:"portalEndpoint,omitempty"`
	Authentication  *AzureStackMetadataAuthentication `json:"authentication,omitempty"`
}

// AzureStackMetadataAuthentication defines configuration for Azure Stack
type AzureStackMetadataAuthentication struct {
	LoginEndpoint string   `json:"loginEndpoint,omitempty"`
	Audiences     []string `json:"audiences,omitempty"`
}

func (az *AzureClient) ValidateDeployment(ctx context.Context, resourceGroupName, deploymentName string, template, params *map[string]interface{}) (valid resources.DeploymentValidateResult, err error) {
	return az.deploymentsClient.Validate(ctx,
		resourceGroupName,
		deploymentName,
		resources.Deployment{
			Properties: &resources.DeploymentProperties{
				Template:   template,
				Parameters: params,
				Mode:       resources.Incremental,
			},
		})
}

func (az *AzureClient) DeployTemplate(ctx context.Context, resourceGroupName, deploymentName string, template, parameters *map[string]interface{}) (de resources.DeploymentExtended, err error) {
	future, err := az.deploymentsClient.CreateOrUpdate(
		ctx,
		resourceGroupName,
		deploymentName,
		resources.Deployment{
			Properties: &resources.DeploymentProperties{
				Template:   template,
				Parameters: parameters,
				Mode:       resources.Incremental,
			},
		})
	if err != nil {
		return de, fmt.Errorf("cannot create deployment: %v", err)
	}

	err = future.WaitForCompletionRef(ctx, az.deploymentsClient.Client)
	if err != nil {
		return de, fmt.Errorf("cannot get the create deployment future response: %v", err)
	}

	return future.Result(az.deploymentsClient)
}

func (az *AzureClient) EnsureResourceGroup(ctx context.Context, name, location string, managedBy *string) (resourceGroup *resources.Group, err error) {
	var tags map[string]*string
	group, err := az.groupsClient.Get(ctx, name)
	if err == nil {
		tags = group.Tags
	}

	response, err := az.groupsClient.CreateOrUpdate(ctx, name, resources.Group{
		Name:      &name,
		Location:  &location,
		ManagedBy: managedBy,
		Tags:      tags,
	})
	if err != nil {
		return &response, err
	}

	return &response, nil
}

func (az *AzureClient) DeleteResourceGroup(ctx context.Context, groupName string) error {
	_, err := az.groupsClient.Get(ctx, groupName)
	if err == nil {
		future, err := az.groupsClient.Delete(ctx, groupName)
		if err != nil {
			return fmt.Errorf("cannot delete resource group %v: %v", groupName, err)
		}
		err = future.WaitForCompletionRef(ctx, az.groupsClient.Client)
		if err != nil {
			// Skip the teardown errors because of https://github.com/Azure/go-autorest/issues/357
			// TODO(feiskyer): fix the issue by upgrading go-autorest version >= v11.3.2.
			log.Printf("Warning: failed to delete resource group %q with error %v", groupName, err)
		}
	}
	return nil
}

func getOAuthConfig(env azure.Environment, subscriptionID, tenantID string) (*adal.OAuthConfig, error) {

	oauthConfig, err := adal.NewOAuthConfig(env.ActiveDirectoryEndpoint, tenantID)
	if err != nil {
		return nil, err
	}

	return oauthConfig, nil
}

func getAzureClient(env azure.Environment, subscriptionID, clientID, tenantID, clientSecret string) (*AzureClient, error) {
	oauthConfig, err := getOAuthConfig(env, subscriptionID, tenantID)
	if err != nil {
		return nil, err
	}

	armSpt, err := adal.NewServicePrincipalToken(*oauthConfig, clientID, clientSecret, env.ServiceManagementEndpoint)
	if err != nil {
		return nil, err
	}

	return getClient(env, subscriptionID, tenantID, armSpt), nil
}

func getClient(env azure.Environment, subscriptionID, tenantID string, armSpt *adal.ServicePrincipalToken) *AzureClient {
	c := &AzureClient{
		environment:    env,
		subscriptionID: subscriptionID,

		deploymentsClient: resources.NewDeploymentsClientWithBaseURI(env.ResourceManagerEndpoint, subscriptionID),
		groupsClient:      resources.NewGroupsClientWithBaseURI(env.ResourceManagerEndpoint, subscriptionID),
	}

	authorizer := autorest.NewBearerAuthorizer(armSpt)
	c.deploymentsClient.Authorizer = authorizer
	c.deploymentsClient.PollingDuration = 60 * time.Minute
	c.groupsClient.Authorizer = authorizer

	return c
}
<|endoftext|>


<commit_before>// +build linux

package leaktest

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"sync"
	"testing"
	"time"
)

// goroutine leaktest code from:
// https://github.com/fortytw2/leaktest
// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
type testReporter struct {
	failed bool
	msg    string
}

func (tr *testReporter) Errorf(format string, args ...interface{}) {
	tr.failed = true
	tr.msg = fmt.Sprintf(format, args)
}

func (tr *testReporter) Fatal(...interface{}) {
}

func tesFn(t *testing.T, inGorutine bool, shouldFail bool, fn func()) {
	checker := &testReporter{}
	checkFn := Check(checker)

	if inGorutine {
		go fn()
	} else {
		fn()
	}

	checkFn()

	if checker.failed != shouldFail {
		if shouldFail {
			t.Errorf("failed to detect leak")
		} else {
			t.Errorf("failed when there was no leak (%s)", checker.msg)
		}
	}
}

func testGoroutineLeak(t *testing.T, fn func()) {
	tesFn(t, true, true, fn)
}

func testLeak(t *testing.T, fn func()) {
	tesFn(t, false, true, fn)
}

func testNoLeak(t *testing.T, fn func()) {
	tesFn(t, false, false, fn)
}

func TestInfiniteForLoop(t *testing.T) {
	testGoroutineLeak(t, func() {
		for {
			time.Sleep(time.Second)
		}
	})
}

func TestSelectUnreferencedChannel(t *testing.T) {
	testGoroutineLeak(t, func() {
		c := make(chan struct{}, 0)
		select {
		case <-c:
		}
	})
}

func TestBlockSelectUnreferencedChannel(t *testing.T) {
	testGoroutineLeak(t, func() {
		c := make(chan struct{}, 0)
		c2 := make(chan struct{}, 0)
		select {
		case <-c:
		case c2 <- struct{}{}:
		}
	})
}

func TestBlockWaitMutexUnreferenced(t *testing.T) {
	testGoroutineLeak(t, func() {
		var mu sync.Mutex
		mu.Lock()
		mu.Lock()
	})
}

func TestBlockWaitRWMutexUnreferenced(t *testing.T) {
	testGoroutineLeak(t, func() {
		var mu sync.RWMutex
		mu.RLock()
		mu.Lock()
	})
}

func TestBlockWaitMutexCondUnreferenced(t *testing.T) {
	testGoroutineLeak(t, func() {
		var mu sync.Mutex
		mu.Lock()
		c := sync.NewCond(&mu)
		c.Wait()
	})
}

func TestLeakFd(t *testing.T) {
	var r, w *os.File
	testLeak(t, func() {
		r, w, _ = os.Pipe()
	})
	r.Close()
	w.Close()
}

func TestLeakFdBeforeClosed(t *testing.T) {
	r, w, _ := os.Pipe()
	var r2, w2 *os.File
	testLeak(t, func() {
		r2, w2, _ = os.Pipe()
		r.Close()
		w.Close()
	})
	r2.Close()
	w2.Close()
}

func TestLeakFdBeforeClosedNoLeak(t *testing.T) {
	r, w, _ := os.Pipe()
	testNoLeak(t, func() {
		r.Close()
		w.Close()
	})
}

func TestLeakChildProcessBeforeClosed(t *testing.T) {
	cmd1 := exec.Command("cat")
	cmd1.Start()
	testLeak(t, func() {
		cmd1.Process.Kill()
		cmd1.Wait()
		cmd2 := exec.Command("cat")
		cmd2.Start()
	})
}

func TestLeakChildProcessBeforeNoLeak(t *testing.T) {
	cmd1 := exec.Command("cat")
	cmd1.Start()
	testNoLeak(t, func() {
		cmd1.Process.Kill()
		cmd1.Wait()
	})
}

func TestLeakChildProcess(t *testing.T) {
	testLeak(t, func() {
		cmd := exec.Command("cat")
		cmd.Start()
	})
}

func TestLeakTempFile(t *testing.T) {
	testLeak(t, func() {
		ioutil.TempFile("", "testleak")
	})
}

func TestEmptyLeak(t *testing.T) {
	defer Check(t)()
	time.Sleep(time.Second)
}
<commit_msg>Cleanup redundant zero arg<commit_after>// +build linux

package leaktest

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"sync"
	"testing"
	"time"
)

// goroutine leaktest code from:
// https://github.com/fortytw2/leaktest
// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
type testReporter struct {
	failed bool
	msg    string
}

func (tr *testReporter) Errorf(format string, args ...interface{}) {
	tr.failed = true
	tr.msg = fmt.Sprintf(format, args)
}

func (tr *testReporter) Fatal(...interface{}) {
}

func tesFn(t *testing.T, inGorutine bool, shouldFail bool, fn func()) {
	checker := &testReporter{}
	checkFn := Check(checker)

	if inGorutine {
		go fn()
	} else {
		fn()
	}

	checkFn()

	if checker.failed != shouldFail {
		if shouldFail {
			t.Errorf("failed to detect leak")
		} else {
			t.Errorf("failed when there was no leak (%s)", checker.msg)
		}
	}
}

func testGoroutineLeak(t *testing.T, fn func()) {
	tesFn(t, true, true, fn)
}

func testLeak(t *testing.T, fn func()) {
	tesFn(t, false, true, fn)
}

func testNoLeak(t *testing.T, fn func()) {
	tesFn(t, false, false, fn)
}

func TestInfiniteForLoop(t *testing.T) {
	testGoroutineLeak(t, func() {
		for {
			time.Sleep(time.Second)
		}
	})
}

func TestSelectUnreferencedChannel(t *testing.T) {
	testGoroutineLeak(t, func() {
		c := make(chan struct{})
		select {
		case <-c:
		}
	})
}

func TestBlockSelectUnreferencedChannel(t *testing.T) {
	testGoroutineLeak(t, func() {
		c := make(chan struct{})
		c2 := make(chan struct{})
		select {
		case <-c:
		case c2 <- struct{}{}:
		}
	})
}

func TestBlockWaitMutexUnreferenced(t *testing.T) {
	testGoroutineLeak(t, func() {
		var mu sync.Mutex
		mu.Lock()
		mu.Lock()
	})
}

func TestBlockWaitRWMutexUnreferenced(t *testing.T) {
	testGoroutineLeak(t, func() {
		var mu sync.RWMutex
		mu.RLock()
		mu.Lock()
	})
}

func TestBlockWaitMutexCondUnreferenced(t *testing.T) {
	testGoroutineLeak(t, func() {
		var mu sync.Mutex
		mu.Lock()
		c := sync.NewCond(&mu)
		c.Wait()
	})
}

func TestLeakFd(t *testing.T) {
	var r, w *os.File
	testLeak(t, func() {
		r, w, _ = os.Pipe()
	})
	r.Close()
	w.Close()
}

func TestLeakFdBeforeClosed(t *testing.T) {
	r, w, _ := os.Pipe()
	var r2, w2 *os.File
	testLeak(t, func() {
		r2, w2, _ = os.Pipe()
		r.Close()
		w.Close()
	})
	r2.Close()
	w2.Close()
}

func TestLeakFdBeforeClosedNoLeak(t *testing.T) {
	r, w, _ := os.Pipe()
	testNoLeak(t, func() {
		r.Close()
		w.Close()
	})
}

func TestLeakChildProcessBeforeClosed(t *testing.T) {
	cmd1 := exec.Command("cat")
	cmd1.Start()
	testLeak(t, func() {
		cmd1.Process.Kill()
		cmd1.Wait()
		cmd2 := exec.Command("cat")
		cmd2.Start()
	})
}

func TestLeakChildProcessBeforeNoLeak(t *testing.T) {
	cmd1 := exec.Command("cat")
	cmd1.Start()
	testNoLeak(t, func() {
		cmd1.Process.Kill()
		cmd1.Wait()
	})
}

func TestLeakChildProcess(t *testing.T) {
	testLeak(t, func() {
		cmd := exec.Command("cat")
		cmd.Start()
	})
}

func TestLeakTempFile(t *testing.T) {
	testLeak(t, func() {
		ioutil.TempFile("", "testleak")
	})
}

func TestEmptyLeak(t *testing.T) {
	defer Check(t)()
	time.Sleep(time.Second)
}
<|endoftext|>


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

package lease

import (
	"fmt"
	"io"
	"sort"

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

// Create a read proxy consisting of the contents defined by the supplied
// refreshers concatenated. See NewReadProxy for more.
//
// If rl is non-nil, it will be used as the first temporary copy of the
// contents, and must match the concatenation of the content returned by the
// refreshers.
func NewMultiReadProxy(
	fl FileLeaser,
	refreshers []Refresher,
	rl ReadLease) (rp ReadProxy) {
	// Create one wrapped read proxy per refresher.
	var wrappedProxies []readProxyAndOffset
	var size int64
	for _, r := range refreshers {
		wrapped := NewReadProxy(fl, r, nil)
		wrappedProxies = append(wrappedProxies, readProxyAndOffset{size, wrapped})
		size += wrapped.Size()
	}

	rp = &multiReadProxy{
		size:  size,
		rps:   wrappedProxies,
		lease: rl,
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////

type multiReadProxy struct {
	// The size of the proxied content.
	size int64

	// The wrapped read proxies, indexed by their logical starting offset.
	//
	// INVARIANT: If len(rps) != 0, rps[0].off == 0
	// INVARIANT: For each x, x.rp.Size() >= 0
	// INVARIANT: For each i>0, rps[i].off == rps[i-i].off + rps[i-i].rp.Size()
	// INVARIANT: size is the sum over the wrapped proxy sizes.
	rps []readProxyAndOffset

	// A read lease for the entire contents. May be nil.
	//
	// INVARIANT: If lease != nil, size == lease.Size()
	lease ReadLease

	destroyed bool
}

func (mrp *multiReadProxy) Size() (size int64) {
	size = mrp.size
	return
}

func (mrp *multiReadProxy) ReadAt(
	ctx context.Context,
	p []byte,
	off int64) (n int, err error) {
	// Special case: we don't support negative offsets, silly user.
	if off < 0 {
		err = fmt.Errorf("Invalid offset: %v", off)
		return
	}

	// Special case: offsets at or beyond the end of our content can never yield
	// any content, and the io.ReaderAt spec allows us to return EOF. Knock them
	// out here so we know off is in range when we start below.
	if off >= mrp.Size() {
		err = io.EOF
		return
	}

	// The read proxy that contains off is the *last* read proxy whose start
	// offset is less than or equal to off. Find the first that is greater and
	// move back one.
	//
	// Because we handled the special cases above, this must be in range.
	wrappedIndex := mrp.upperBound(off) - 1

	if wrappedIndex < 0 || wrappedIndex >= len(mrp.rps) {
		panic(fmt.Sprintf("Unexpected index: %v", wrappedIndex))
	}

	// Keep going until we've got nothing left to do.
	for len(p) > 0 {
		// Have we run out of wrapped read proxies?
		if wrappedIndex == len(mrp.rps) {
			err = io.EOF
			return
		}

		// Read from the wrapped proxy, accumulating into our total before checking
		// for a read error.
		wrappedN, wrappedErr := mrp.readFromOne(ctx, wrappedIndex, p, off)
		n += wrappedN
		if wrappedErr != nil {
			err = wrappedErr
			return
		}

		// readFromOne guarantees to either fill our buffer or exhaust the wrapped
		// proxy. So advance the buffer, the offset, and the wrapped proxy index
		// and go again.
		p = p[wrappedN:]
		off += int64(wrappedN)
		wrappedIndex++
	}

	return
}

func (mrp *multiReadProxy) Upgrade(
	ctx context.Context) (rwl ReadWriteLease, err error) {
	panic("TODO")
}

func (mrp *multiReadProxy) Destroy() {
	// Destroy all of the wrapped proxies.
	for _, entry := range mrp.rps {
		entry.rp.Destroy()
	}

	// Destroy the lease for the entire contents, if any.
	if mrp.lease != nil {
		mrp.lease.Revoke()
	}

	// Crash early if called again.
	mrp.rps = nil
	mrp.lease = nil
}

func (mrp *multiReadProxy) CheckInvariants() {
	// INVARIANT: If len(rps) != 0, rps[0].off == 0
	if len(mrp.rps) != 0 && mrp.rps[0].off != 0 {
		panic(fmt.Sprintf("Unexpected starting point: %v", mrp.rps[0].off))
	}

	// INVARIANT: For each x, x.rp.Size() >= 0
	for _, x := range mrp.rps {
		if x.rp.Size() < 0 {
			panic(fmt.Sprintf("Negative size: %v", x.rp.Size()))
		}
	}

	// INVARIANT: For each i>0, rps[i].off == rps[i-i].off + rps[i-i].rp.Size()
	for i := range mrp.rps {
		if i > 0 && !(mrp.rps[i].off == mrp.rps[i-1].off+mrp.rps[i-1].rp.Size()) {
			panic("Offsets are not indexed correctly.")
		}
	}

	// INVARIANT: size is the sum over the wrapped proxy sizes.
	var sum int64
	for _, wrapped := range mrp.rps {
		sum += wrapped.rp.Size()
	}

	if sum != mrp.size {
		panic(fmt.Sprintf("Size mismatch: %v vs. %v", sum, mrp.size))
	}

	// INVARIANT: If lease != nil, size == lease.Size()
	if mrp.lease != nil && mrp.size != mrp.lease.Size() {
		panic(fmt.Sprintf("Size mismatch: %v vs. %v", mrp.size, mrp.lease.Size()))
	}
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type readProxyAndOffset struct {
	off int64
	rp  ReadProxy
}

// Return the index within mrp.rps of the first read proxy whose logical offset
// is greater than off. If there is none, return len(mrp.rps).
func (mrp *multiReadProxy) upperBound(off int64) (index int) {
	pred := func(i int) bool {
		return mrp.rps[i].off > off
	}

	return sort.Search(len(mrp.rps), pred)
}

// Serve a read from the wrapped proxy at the given index within our array of
// wrapped proxies. The offset is relative to the start of the multiReadProxy,
// not the wrapped proxy.
//
// Guarantees, letting wrapped be mrp.rps[i].rp:
//
//  *  If err == nil, either n == len(p) or n == wrapped.Size().
//  *  Never returns err == io.EOF.
//
// REQUIRES: index < len(mrp.rps)
// REQUIRES: mrp.rps[index].off <= off < mrp.rps[index].off + wrapped.Size()
func (mrp *multiReadProxy) readFromOne(
	ctx context.Context,
	index int,
	p []byte,
	off int64) (n int, err error) {
	// Check input requirements.
	panic("TODO")

	// Check guarantees on return.
	defer func() {
		panic("TODO")
	}()

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

package lease

import (
	"fmt"
	"io"
	"sort"

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

// Create a read proxy consisting of the contents defined by the supplied
// refreshers concatenated. See NewReadProxy for more.
//
// If rl is non-nil, it will be used as the first temporary copy of the
// contents, and must match the concatenation of the content returned by the
// refreshers.
func NewMultiReadProxy(
	fl FileLeaser,
	refreshers []Refresher,
	rl ReadLease) (rp ReadProxy) {
	// Create one wrapped read proxy per refresher.
	var wrappedProxies []readProxyAndOffset
	var size int64
	for _, r := range refreshers {
		wrapped := NewReadProxy(fl, r, nil)
		wrappedProxies = append(wrappedProxies, readProxyAndOffset{size, wrapped})
		size += wrapped.Size()
	}

	rp = &multiReadProxy{
		size:  size,
		rps:   wrappedProxies,
		lease: rl,
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////

type multiReadProxy struct {
	// The size of the proxied content.
	size int64

	// The wrapped read proxies, indexed by their logical starting offset.
	//
	// INVARIANT: If len(rps) != 0, rps[0].off == 0
	// INVARIANT: For each x, x.rp.Size() >= 0
	// INVARIANT: For each i>0, rps[i].off == rps[i-i].off + rps[i-i].rp.Size()
	// INVARIANT: size is the sum over the wrapped proxy sizes.
	rps []readProxyAndOffset

	// A read lease for the entire contents. May be nil.
	//
	// INVARIANT: If lease != nil, size == lease.Size()
	lease ReadLease

	destroyed bool
}

func (mrp *multiReadProxy) Size() (size int64) {
	size = mrp.size
	return
}

func (mrp *multiReadProxy) ReadAt(
	ctx context.Context,
	p []byte,
	off int64) (n int, err error) {
	// Special case: we don't support negative offsets, silly user.
	if off < 0 {
		err = fmt.Errorf("Invalid offset: %v", off)
		return
	}

	// Special case: offsets at or beyond the end of our content can never yield
	// any content, and the io.ReaderAt spec allows us to return EOF. Knock them
	// out here so we know off is in range when we start below.
	if off >= mrp.Size() {
		err = io.EOF
		return
	}

	// The read proxy that contains off is the *last* read proxy whose start
	// offset is less than or equal to off. Find the first that is greater and
	// move back one.
	//
	// Because we handled the special cases above, this must be in range.
	wrappedIndex := mrp.upperBound(off) - 1

	if wrappedIndex < 0 || wrappedIndex >= len(mrp.rps) {
		panic(fmt.Sprintf("Unexpected index: %v", wrappedIndex))
	}

	// Keep going until we've got nothing left to do.
	for len(p) > 0 {
		// Have we run out of wrapped read proxies?
		if wrappedIndex == len(mrp.rps) {
			err = io.EOF
			return
		}

		// Read from the wrapped proxy, accumulating into our total before checking
		// for a read error.
		wrappedN, wrappedErr := mrp.readFromOne(ctx, wrappedIndex, p, off)
		n += wrappedN
		if wrappedErr != nil {
			err = wrappedErr
			return
		}

		// readFromOne guarantees to either fill our buffer or exhaust the wrapped
		// proxy. So advance the buffer, the offset, and the wrapped proxy index
		// and go again.
		p = p[wrappedN:]
		off += int64(wrappedN)
		wrappedIndex++
	}

	return
}

func (mrp *multiReadProxy) Upgrade(
	ctx context.Context) (rwl ReadWriteLease, err error) {
	panic("TODO")
}

func (mrp *multiReadProxy) Destroy() {
	// Destroy all of the wrapped proxies.
	for _, entry := range mrp.rps {
		entry.rp.Destroy()
	}

	// Destroy the lease for the entire contents, if any.
	if mrp.lease != nil {
		mrp.lease.Revoke()
	}

	// Crash early if called again.
	mrp.rps = nil
	mrp.lease = nil
}

func (mrp *multiReadProxy) CheckInvariants() {
	// INVARIANT: If len(rps) != 0, rps[0].off == 0
	if len(mrp.rps) != 0 && mrp.rps[0].off != 0 {
		panic(fmt.Sprintf("Unexpected starting point: %v", mrp.rps[0].off))
	}

	// INVARIANT: For each x, x.rp.Size() >= 0
	for _, x := range mrp.rps {
		if x.rp.Size() < 0 {
			panic(fmt.Sprintf("Negative size: %v", x.rp.Size()))
		}
	}

	// INVARIANT: For each i>0, rps[i].off == rps[i-i].off + rps[i-i].rp.Size()
	for i := range mrp.rps {
		if i > 0 && !(mrp.rps[i].off == mrp.rps[i-1].off+mrp.rps[i-1].rp.Size()) {
			panic("Offsets are not indexed correctly.")
		}
	}

	// INVARIANT: size is the sum over the wrapped proxy sizes.
	var sum int64
	for _, wrapped := range mrp.rps {
		sum += wrapped.rp.Size()
	}

	if sum != mrp.size {
		panic(fmt.Sprintf("Size mismatch: %v vs. %v", sum, mrp.size))
	}

	// INVARIANT: If lease != nil, size == lease.Size()
	if mrp.lease != nil && mrp.size != mrp.lease.Size() {
		panic(fmt.Sprintf("Size mismatch: %v vs. %v", mrp.size, mrp.lease.Size()))
	}
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

type readProxyAndOffset struct {
	off int64
	rp  ReadProxy
}

// Return the index within mrp.rps of the first read proxy whose logical offset
// is greater than off. If there is none, return len(mrp.rps).
func (mrp *multiReadProxy) upperBound(off int64) (index int) {
	pred := func(i int) bool {
		return mrp.rps[i].off > off
	}

	return sort.Search(len(mrp.rps), pred)
}

// Serve a read from the wrapped proxy at the given index within our array of
// wrapped proxies. The offset is relative to the start of the multiReadProxy,
// not the wrapped proxy.
//
// Guarantees, letting wrapped be mrp.rps[i].rp:
//
//  *  If err == nil, either n == len(p) or n == wrapped.Size().
//  *  Never returns err == io.EOF.
//
// REQUIRES: index < len(mrp.rps)
// REQUIRES: mrp.rps[index].off <= off < mrp.rps[index].off + wrapped.Size()
func (mrp *multiReadProxy) readFromOne(
	ctx context.Context,
	index int,
	p []byte,
	off int64) (n int, err error) {
	// Check input requirements.
	if !(index < len(mrp.rps)) {
		panic(fmt.Sprintf("Out of range wrapped index: %v", index))
	}

	wrapped := mrp.rps[index].rp
	wrappedStart := mrp.rps[index].off
	if !(wrappedStart <= off && off < wrappedStart+wrapped.Size()) {
		panic(fmt.Sprintf(
			"Offset %v not in range [%v, %v)",
			off,
			wrappedStart,
			wrappedStart+wrapped.Size()))
	}

	// Check guarantees on return.
	defer func() {
		panic("TODO")
	}()

	panic("TODO")
}
<|endoftext|>


<commit_before>// Copyright 2014 Brave New Software

//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at

//        http://www.apache.org/licenses/LICENSE-2.0

//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.
//

package statshub

import (
	"github.com/garyburd/redigo/redis"
	"log"
	"os"
	"time"
)

const (
	redisConnectTimeout = 10 * time.Second
	redisReadTimeout    = 10 * time.Second
	redisWriteTimeout   = 10 * time.Second
)

var (
	pool *redis.Pool
)

// init() initializes our redis environment by setting up a connection pool
func init() {
	pool = &redis.Pool{
		MaxIdle:     100,
		MaxActive:   1000,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.DialTimeout(
				"tcp",
				os.Getenv("REDIS_ADDR"),
				redisConnectTimeout,
				redisReadTimeout,
				redisWriteTimeout)
			if err != nil {
				log.Fatalf("Unable to dial redis: %s", err)
			}
			if _, err := c.Do("AUTH", os.Getenv("REDIS_PASS")); err != nil {
				c.Close()
				log.Fatalf("Unable to authenticate to redis: %s", err)
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

// redisConn is a wrapper for a redis.Conn that itself implements the
// redis.Conn interface. Unlike a normal redis.Conn, redisConn stops processing
// new commands after it encountersits first error.
type redisConn struct {
	orig redis.Conn
	err  error
}

// connectToRedis() connects to our cloud Redis server and authenticates
func connectToRedis() (conn redis.Conn, err error) {
	return &redisConn{orig: pool.Get()}, nil
}

func (conn *redisConn) Close() (err error) {
	return conn.orig.Close()
}

func (conn *redisConn) Err() error {
	return conn.err
}

func (conn *redisConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
	if conn.err != nil {
		return nil, conn.err
	} else {
		reply, err = conn.orig.Do(commandName, args...)
		conn.err = err
		return
	}
}

func (conn *redisConn) Send(commandName string, args ...interface{}) (err error) {
	if conn.err != nil {
		return conn.err
	} else {
		err = conn.orig.Send(commandName, args...)
		conn.err = err
		return
	}
}

func (conn *redisConn) Flush() (err error) {
	if conn.err != nil {
		return conn.err
	} else {
		err = conn.orig.Flush()
		conn.err = err
		return
	}
}

func (conn *redisConn) Receive() (reply interface{}, err error) {
	if conn.err != nil {
		return nil, conn.err
	} else {
		reply, err = conn.orig.Receive()
		conn.err = err
		return
	}
}
<commit_msg>#4 Not terminating just because I couldn't dial Redis<commit_after>// Copyright 2014 Brave New Software

//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at

//        http://www.apache.org/licenses/LICENSE-2.0

//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.
//

package statshub

import (
	"fmt"
	"github.com/garyburd/redigo/redis"
	"os"
	"time"
)

const (
	redisConnectTimeout = 10 * time.Second
	redisReadTimeout    = 10 * time.Second
	redisWriteTimeout   = 10 * time.Second
)

var (
	pool *redis.Pool
)

// init() initializes our redis environment by setting up a connection pool
func init() {
	pool = &redis.Pool{
		MaxIdle:     100,
		MaxActive:   1000,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.DialTimeout(
				"tcp",
				os.Getenv("REDIS_ADDR"),
				redisConnectTimeout,
				redisReadTimeout,
				redisWriteTimeout)

			if err != nil {
				return nil, fmt.Errorf("Unable to dial redis: %s", err)
			}
			if _, err := c.Do("AUTH", os.Getenv("REDIS_PASS")); err != nil {
				c.Close()
				return nil, fmt.Errorf("Unable to authenticate to redis: %s", err)
			}
			return c, err
		},
		TestOnBorrow: func(c redis.Conn, t time.Time) error {
			_, err := c.Do("PING")
			return err
		},
	}
}

// redisConn is a wrapper for a redis.Conn that itself implements the
// redis.Conn interface. Unlike a normal redis.Conn, redisConn stops processing
// new commands after it encountersits first error.
type redisConn struct {
	orig redis.Conn
	err  error
}

// connectToRedis() connects to our cloud Redis server and authenticates
func connectToRedis() (conn redis.Conn, err error) {
	return &redisConn{orig: pool.Get()}, nil
}

func (conn *redisConn) Close() (err error) {
	return conn.orig.Close()
}

func (conn *redisConn) Err() error {
	return conn.err
}

func (conn *redisConn) Do(commandName string, args ...interface{}) (reply interface{}, err error) {
	if conn.err != nil {
		return nil, conn.err
	} else {
		reply, err = conn.orig.Do(commandName, args...)
		conn.err = err
		return
	}
}

func (conn *redisConn) Send(commandName string, args ...interface{}) (err error) {
	if conn.err != nil {
		return conn.err
	} else {
		err = conn.orig.Send(commandName, args...)
		conn.err = err
		return
	}
}

func (conn *redisConn) Flush() (err error) {
	if conn.err != nil {
		return conn.err
	} else {
		err = conn.orig.Flush()
		conn.err = err
		return
	}
}

func (conn *redisConn) Receive() (reply interface{}, err error) {
	if conn.err != nil {
		return nil, conn.err
	} else {
		reply, err = conn.orig.Receive()
		conn.err = err
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"net/url"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

type natsMetrics struct {
	Connections float64 `json:"connections"`
	Routes      float64 `json:"routes"`

	MessagesIn  float64 `json:"in_msgs"`
	MessagesOut float64 `json:"out_msgs"`

	BytesIn  float64 `json:"in_bytes"`
	BytesOut float64 `json:"out_bytes"`

	SlowConsumers float64 `json:"slow_consumers"`
}

var (
	connections = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "connections",
		Help:      "Active connections to gnatsd",
	})

	routes = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "routes",
		Help:      "Active routes to gnatsd",
	})

	messageCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gnatsd",
		Name:      "msg_total",
		Help:      "Count of transferred messages",
	},
		[]string{"direction"},
	)

	bytesCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gnatsd",
		Name:      "bytes_total",
		Help:      "Amount of transmitted data",
	},
		[]string{"direction"},
	)

	slowConsumers = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "slow_consumers",
		Help:      "Number of slow consumers",
	})
)

// Exporter implements the prometheus.Collector interface. It exposes the metrics
// of a NATS node.
type Exporter struct {
	NatsURL string
}

// NewExporter instantiates a new NATS Exporter.
func NewExporter(natsURL *url.URL) *Exporter {
	natsURL.Path = "/varz"
	return &Exporter{
		NatsURL: natsURL.String(),
	}
}

// Describe describes all the registered stats metrics from the NATS node.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	connections.Describe(ch)
	routes.Describe(ch)
	messageCounter.Describe(ch)
	bytesCounter.Describe(ch)
	slowConsumers.Describe(ch)
}

// Collect collects all the registered stats metrics from the NATS node.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.collect()

	connections.Collect(ch)
	routes.Collect(ch)
	messageCounter.Collect(ch)
	bytesCounter.Collect(ch)
	slowConsumers.Collect(ch)
}

func (e *Exporter) collect() {
	var metrics natsMetrics

	httpClient := http.DefaultClient
	httpClient.Timeout = 1 * time.Second
	resp, err := httpClient.Get(e.NatsURL)
	if err != nil {
		log.Printf("could not retrieve NATS metrics: %v", err)
		return
	}

	err = json.NewDecoder(resp.Body).Decode(&metrics)
	if err != nil {
		log.Printf("could not decode NATS metrics: %v", err)
	}

	connections.Set(metrics.Connections)
	routes.Set(metrics.Routes)

	messageCounter.WithLabelValues("in").Add(metrics.MessagesIn)
	messageCounter.WithLabelValues("out").Add(metrics.MessagesOut)

	bytesCounter.WithLabelValues("in").Add(metrics.BytesIn)
	bytesCounter.WithLabelValues("out").Add(metrics.BytesOut)

	slowConsumers.Set(metrics.SlowConsumers)
}
<commit_msg>fixing metric generation by replacing Add with Set<commit_after>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"net/url"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

type natsMetrics struct {
	Connections float64 `json:"connections"`
	Routes      float64 `json:"routes"`

	MessagesIn  float64 `json:"in_msgs"`
	MessagesOut float64 `json:"out_msgs"`

	BytesIn  float64 `json:"in_bytes"`
	BytesOut float64 `json:"out_bytes"`

	SlowConsumers float64 `json:"slow_consumers"`
}

var (
	connections = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "connections",
		Help:      "Active connections to gnatsd",
	})

	routes = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "routes",
		Help:      "Active routes to gnatsd",
	})

	messageCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gnatsd",
		Name:      "msg_total",
		Help:      "Count of transferred messages",
	},
		[]string{"direction"},
	)

	bytesCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "gnatsd",
		Name:      "bytes_total",
		Help:      "Amount of transmitted data",
	},
		[]string{"direction"},
	)

	slowConsumers = prometheus.NewGauge(prometheus.GaugeOpts{
		Namespace: "gnatsd",
		Name:      "slow_consumers",
		Help:      "Number of slow consumers",
	})
)

// Exporter implements the prometheus.Collector interface. It exposes the metrics
// of a NATS node.
type Exporter struct {
	NatsURL string
}

// NewExporter instantiates a new NATS Exporter.
func NewExporter(natsURL *url.URL) *Exporter {
	natsURL.Path = "/varz"
	return &Exporter{
		NatsURL: natsURL.String(),
	}
}

// Describe describes all the registered stats metrics from the NATS node.
func (e *Exporter) Describe(ch chan<- *prometheus.Desc) {
	connections.Describe(ch)
	routes.Describe(ch)
	messageCounter.Describe(ch)
	bytesCounter.Describe(ch)
	slowConsumers.Describe(ch)
}

// Collect collects all the registered stats metrics from the NATS node.
func (e *Exporter) Collect(ch chan<- prometheus.Metric) {
	e.collect()

	connections.Collect(ch)
	routes.Collect(ch)
	messageCounter.Collect(ch)
	bytesCounter.Collect(ch)
	slowConsumers.Collect(ch)
}

func (e *Exporter) collect() {
	var metrics natsMetrics

	httpClient := http.DefaultClient
	httpClient.Timeout = 1 * time.Second
	resp, err := httpClient.Get(e.NatsURL)
	if err != nil {
		log.Printf("could not retrieve NATS metrics: %v", err)
		return
	}

	err = json.NewDecoder(resp.Body).Decode(&metrics)
	if err != nil {
		log.Printf("could not decode NATS metrics: %v", err)
	}

	connections.Set(metrics.Connections)
	routes.Set(metrics.Routes)

	messageCounter.WithLabelValues("in").Set(metrics.MessagesIn)
	messageCounter.WithLabelValues("out").Set(metrics.MessagesOut)

	bytesCounter.WithLabelValues("in").Set(metrics.BytesIn)
	bytesCounter.WithLabelValues("out").Set(metrics.BytesOut)

	slowConsumers.Set(metrics.SlowConsumers)
}
<|endoftext|>


<commit_before>package utp

import (
	"net"
	"testing"
)

func getTCPConnectionPair() (net.Conn, net.Conn, error) {
	lst, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, nil, err
	}

	var conn0 net.Conn
	var err0 error
	done := make(chan struct{})
	go func() {
		conn0, err0 = lst.Accept()
		close(done)
	}()

	conn1, err := net.Dial("tcp", lst.Addr().String())
	if err != nil {
		return nil, nil, err
	}

	<-done
	if err0 != nil {
		return nil, nil, err0
	}
	return conn0, conn1, nil
}

func getUTPConnectionPair() (net.Conn, net.Conn, error) {
	lst, err := NewSocket("udp", "127.0.0.1:0")
	if err != nil {
		return nil, nil, err
	}
	defer lst.Close()

	var conn0 net.Conn
	var err0 error
	done := make(chan struct{})
	go func() {
		conn0, err0 = lst.Accept()
		close(done)
	}()

	conn1, err := Dial(lst.Addr().String())
	if err != nil {
		return nil, nil, err
	}

	<-done
	if err0 != nil {
		return nil, nil, err0
	}

	return conn0, conn1, nil
}

func benchConnPair(b *testing.B, c0, c1 net.Conn) {
	b.ReportAllocs()
	b.SetBytes(128 << 10)
	b.ResetTimer()

	request := make([]byte, 52)
	response := make([]byte, (128<<10)+8)

	pair := []net.Conn{c0, c1}
	for i := 0; i < b.N; i++ {
		if i%2 == 0 {
			pair[0] = c0
			pair[1] = c1
		} else {
			pair[0] = c1
			pair[1] = c0
		}

		if _, err := pair[0].Write(request); err != nil {
			b.Fatal(err)
		}

		if _, err := pair[1].Read(request[:8]); err != nil {
			b.Fatal(err)
		}
		if _, err := pair[1].Read(request[8:]); err != nil {
			b.Fatal(err)
		}
		if _, err := pair[1].Write(response); err != nil {
			b.Fatal(err)
		}
		if _, err := pair[0].Read(response[:8]); err != nil {
			b.Fatal(err)
		}
		if _, err := pair[0].Read(response[8:]); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkSyncthingTCP(b *testing.B) {
	conn0, conn1, err := getTCPConnectionPair()
	if err != nil {
		b.Fatal(err)
	}

	defer conn0.Close()
	defer conn1.Close()

	benchConnPair(b, conn0, conn1)
}

func BenchmarkSyncthingUDPUTP(b *testing.B) {
	conn0, conn1, err := getUTPConnectionPair()
	if err != nil {
		b.Fatal(err)
	}

	defer conn0.Close()
	defer conn1.Close()

	benchConnPair(b, conn0, conn1)
}

func BenchmarkSyncthingInprocUTP(b *testing.B) {
	c0, c1 := connPair()
	defer c0.Close()
	defer c1.Close()
	benchConnPair(b, c0, c1)
}
<commit_msg>Fix the syncthing tests to do read and write concurrently<commit_after>package utp

import (
	"io"
	"net"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"
)

func getTCPConnectionPair() (net.Conn, net.Conn, error) {
	lst, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		return nil, nil, err
	}

	var conn0 net.Conn
	var err0 error
	done := make(chan struct{})
	go func() {
		conn0, err0 = lst.Accept()
		close(done)
	}()

	conn1, err := net.Dial("tcp", lst.Addr().String())
	if err != nil {
		return nil, nil, err
	}

	<-done
	if err0 != nil {
		return nil, nil, err0
	}
	return conn0, conn1, nil
}

func getUTPConnectionPair() (net.Conn, net.Conn, error) {
	lst, err := NewSocket("udp", "127.0.0.1:0")
	if err != nil {
		return nil, nil, err
	}
	defer lst.Close()

	var conn0 net.Conn
	var err0 error
	done := make(chan struct{})
	go func() {
		conn0, err0 = lst.Accept()
		close(done)
	}()

	conn1, err := Dial(lst.Addr().String())
	if err != nil {
		return nil, nil, err
	}

	<-done
	if err0 != nil {
		return nil, nil, err0
	}

	return conn0, conn1, nil
}

func requireWriteAll(t testing.TB, b []byte, w io.Writer) {
	n, err := w.Write(b)
	require.NoError(t, err)
	require.EqualValues(t, len(b), n)
}

func requireReadExactly(t testing.TB, b []byte, r io.Reader) {
	n, err := io.ReadFull(r, b)
	require.NoError(t, err)
	require.EqualValues(t, len(b), n)
}

func benchConnPair(b *testing.B, c0, c1 net.Conn) {
	b.ReportAllocs()
	request := make([]byte, 52)
	response := make([]byte, (128<<10)+8)
	b.SetBytes(int64(len(request) + len(response)))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		c, s := func() (net.Conn, net.Conn) {
			if i%2 == 0 {
				return c0, c1
			} else {
				return c1, c0
			}
		}()
		var wg sync.WaitGroup
		wg.Add(2)
		go func() {
			defer wg.Done()
			requireWriteAll(b, request, c)
			requireReadExactly(b, response[:8], c)
			requireReadExactly(b, response[8:], c)
		}()
		go func() {
			defer wg.Done()
			requireReadExactly(b, request[:8], s)
			requireReadExactly(b, request[8:], s)
			requireWriteAll(b, response, s)
		}()
		wg.Wait()
	}
}

func BenchmarkSyncthingTCP(b *testing.B) {
	conn0, conn1, err := getTCPConnectionPair()
	if err != nil {
		b.Fatal(err)
	}

	defer conn0.Close()
	defer conn1.Close()

	benchConnPair(b, conn0, conn1)
}

func BenchmarkSyncthingUDPUTP(b *testing.B) {
	conn0, conn1, err := getUTPConnectionPair()
	if err != nil {
		b.Fatal(err)
	}

	defer conn0.Close()
	defer conn1.Close()

	benchConnPair(b, conn0, conn1)
}

func BenchmarkSyncthingInprocUTP(b *testing.B) {
	c0, c1 := connPair()
	defer c0.Close()
	defer c1.Close()
	benchConnPair(b, c0, c1)
}
<|endoftext|>


<commit_before>package fsutil

import (
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path"
	"testing"
	"time"

	"github.com/Symantec/Dominator/lib/log/testlogger"
)

var errorTimeout = errors.New("timeout")

func TestWatchFileDir(t *testing.T) {
	dirname, err := ioutil.TempDir("", "WatchFileTests")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirname)
	testWatchFile(t, dirname)
}

func TestWatchFileCwd(t *testing.T) {
	cwd, err := os.Getwd()
	if err != nil {
		t.Fatal(err)
	}
	defer os.Chdir(cwd)
	dirname, err := ioutil.TempDir("", "WatchFileTests")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirname)
	if err := os.Chdir(dirname); err != nil {
		t.Fatal(err)
	}
	testWatchFile(t, "")
}

func testWatchFile(t *testing.T, dirname string) {
	logger := testlogger.New(t)
	pathNotExist := path.Join(dirname, "never-exists")
	ch := WatchFile(pathNotExist, logger)
	rc, err := watchTimeout(ch, time.Millisecond*10)
	if err != errorTimeout {
		rc.Close()
		t.Fatal("Expected timeout error for non-existant file")
	}
	pathExists := path.Join(dirname, "exists")
	file, err := os.Create(pathExists)
	if err != nil {
		t.Fatal(err)
	}
	file.Close()
	ch = WatchFile(pathExists, logger)
	rc, err = watchTimeout(ch, time.Millisecond*10)
	if err != nil {
		t.Fatal(err)
	} else {
		rc.Close()
	}
	pathExistsLater := path.Join(dirname, "exists-later")
	go func() {
		time.Sleep(time.Millisecond * 50)
		file, err := os.Create(pathExistsLater)
		if err != nil {
			panic(err)
		}
		file.Close()
	}()
	ch = WatchFile(pathExistsLater, logger)
	_, err = watchTimeout(ch, time.Millisecond*10)
	if err != errorTimeout {
		t.Fatal("Expected timeout error for non-existant file")
	}
	rc, err = watchTimeout(ch, time.Millisecond*90)
	if err != nil {
		t.Fatal(err)
	} else {
		rc.Close()
	}
	pathWillBeRenamed := path.Join(dirname, "will-be-renamed")
	file, err = os.Create(pathWillBeRenamed)
	if err != nil {
		t.Fatal(err)
	}
	file.Close()
	rc, err = watchTimeout(ch, time.Millisecond*10)
	if err != errorTimeout {
		rc.Close()
		t.Fatal("Expected timeout error for unchanged file")
	}
	if err := os.Rename(pathWillBeRenamed, pathExistsLater); err != nil {
		t.Fatal(err)
	}
	rc, err = watchTimeout(ch, time.Millisecond*10)
	if err != nil {
		t.Fatal(err)
	} else {
		rc.Close()
	}
}

func watchTimeout(channel <-chan io.ReadCloser, timeout time.Duration) (
	io.ReadCloser, error) {
	select {
	case readCloser := <-channel:
		return readCloser, nil
	case <-time.After(timeout):
		return nil, errorTimeout
	}
	panic("impossible")
}
<commit_msg>Increase lib/fsutil test timeouts to 50 ms to help TravisCI pass.<commit_after>package fsutil

import (
	"errors"
	"io"
	"io/ioutil"
	"os"
	"path"
	"testing"
	"time"

	"github.com/Symantec/Dominator/lib/log/testlogger"
)

var errorTimeout = errors.New("timeout")

func TestWatchFileDir(t *testing.T) {
	dirname, err := ioutil.TempDir("", "WatchFileTests")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirname)
	testWatchFile(t, dirname)
}

func TestWatchFileCwd(t *testing.T) {
	cwd, err := os.Getwd()
	if err != nil {
		t.Fatal(err)
	}
	defer os.Chdir(cwd)
	dirname, err := ioutil.TempDir("", "WatchFileTests")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dirname)
	if err := os.Chdir(dirname); err != nil {
		t.Fatal(err)
	}
	testWatchFile(t, "")
}

func testWatchFile(t *testing.T, dirname string) {
	logger := testlogger.New(t)
	pathNotExist := path.Join(dirname, "never-exists")
	ch := WatchFile(pathNotExist, logger)
	rc, err := watchTimeout(ch, time.Millisecond*50)
	if err != errorTimeout {
		rc.Close()
		t.Fatal("Expected timeout error for non-existant file")
	}
	pathExists := path.Join(dirname, "exists")
	file, err := os.Create(pathExists)
	if err != nil {
		t.Fatal(err)
	}
	file.Close()
	ch = WatchFile(pathExists, logger)
	rc, err = watchTimeout(ch, time.Millisecond*50)
	if err != nil {
		t.Fatal(err)
	} else {
		rc.Close()
	}
	pathExistsLater := path.Join(dirname, "exists-later")
	go func() {
		time.Sleep(time.Millisecond * 50)
		file, err := os.Create(pathExistsLater)
		if err != nil {
			panic(err)
		}
		file.Close()
	}()
	ch = WatchFile(pathExistsLater, logger)
	_, err = watchTimeout(ch, time.Millisecond*10)
	if err != errorTimeout {
		t.Fatal("Expected timeout error for non-existant file")
	}
	rc, err = watchTimeout(ch, time.Millisecond*90)
	if err != nil {
		t.Fatal(err)
	} else {
		rc.Close()
	}
	pathWillBeRenamed := path.Join(dirname, "will-be-renamed")
	file, err = os.Create(pathWillBeRenamed)
	if err != nil {
		t.Fatal(err)
	}
	file.Close()
	rc, err = watchTimeout(ch, time.Millisecond*10)
	if err != errorTimeout {
		rc.Close()
		t.Fatal("Expected timeout error for unchanged file")
	}
	if err := os.Rename(pathWillBeRenamed, pathExistsLater); err != nil {
		t.Fatal(err)
	}
	rc, err = watchTimeout(ch, time.Millisecond*10)
	if err != nil {
		t.Fatal(err)
	} else {
		rc.Close()
	}
}

func watchTimeout(channel <-chan io.ReadCloser, timeout time.Duration) (
	io.ReadCloser, error) {
	select {
	case readCloser := <-channel:
		return readCloser, nil
	case <-time.After(timeout):
		return nil, errorTimeout
	}
	panic("impossible")
}
<|endoftext|>


<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package thrift

import (
	"bytes"
	"io"
	"net/http"
	"net/url"
	"strconv"
)

type THttpClient struct {
	response           *http.Response
	url                *url.URL
	requestBuffer      *bytes.Buffer
	header             http.Header
	nsecConnectTimeout int64
	nsecReadTimeout    int64
}

type THttpClientTransportFactory struct {
	url    string
	isPost bool
}

func (p *THttpClientTransportFactory) GetTransport(trans TTransport) TTransport {
	if trans != nil {
		t, ok := trans.(*THttpClient)
		if ok && t.url != nil {
			if t.requestBuffer != nil {
				t2, _ := NewTHttpPostClient(t.url.String())
				return t2
			}
			t2, _ := NewTHttpClient(t.url.String())
			return t2
		}
	}
	if p.isPost {
		s, _ := NewTHttpPostClient(p.url)
		return s
	}
	s, _ := NewTHttpClient(p.url)
	return s
}

func NewTHttpClientTransportFactory(url string) *THttpClientTransportFactory {
	return &THttpClientTransportFactory{url: url, isPost: false}
}

func NewTHttpPostClientTransportFactory(url string) *THttpClientTransportFactory {
	return &THttpClientTransportFactory{url: url, isPost: true}
}

func NewTHttpClient(urlstr string) (TTransport, error) {
	parsedURL, err := url.Parse(urlstr)
	if err != nil {
		return nil, err
	}
	response, err := http.Get(urlstr)
	if err != nil {
		return nil, err
	}
	return &THttpClient{response: response, url: parsedURL}, nil
}

func NewTHttpPostClient(urlstr string) (TTransport, error) {
	parsedURL, err := url.Parse(urlstr)
	if err != nil {
		return nil, err
	}
	buf := make([]byte, 0, 1024)
	return &THttpClient{url: parsedURL, requestBuffer: bytes.NewBuffer(buf), header: http.Header{}}, nil
}

// Set the HTTP Header for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// httpTrans.SetHeader("User-Agent","Thrift Client 1.0")
func (p *THttpClient) SetHeader(key string, value string) {
	p.header.Add(key, value)
}

// Get the HTTP Header represented by the supplied Header Key for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// hdrValue := httpTrans.GetHeader("User-Agent")
func (p *THttpClient) GetHeader(key string) string {
	return p.header.Get(key)
}

// Deletes the HTTP Header given a Header Key for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// httpTrans.DelHeader("User-Agent")
func (p *THttpClient) DelHeader(key string) {
	p.header.Del(key)
}

func (p *THttpClient) Open() error {
	// do nothing
	return nil
}

func (p *THttpClient) IsOpen() bool {
	return p.response != nil || p.requestBuffer != nil
}

func (p *THttpClient) Close() error {
	if p.response != nil && p.response.Body != nil {
		err := p.response.Body.Close()
		p.response = nil
		return err
	}
	if p.requestBuffer != nil {
		p.requestBuffer.Reset()
		p.requestBuffer = nil
	}
	return nil
}

func (p *THttpClient) Read(buf []byte) (int, error) {
	if p.response == nil {
		return 0, NewTTransportException(NOT_OPEN, "Response buffer is empty, no request.")
	}
	n, err := p.response.Body.Read(buf)
	if n > 0 && (err == nil || err == io.EOF) {
		return n, nil
	}
	return n, NewTTransportExceptionFromError(err)
}

func (p *THttpClient) ReadByte() (c byte, err error) {
	return readByte(p.response.Body)
}

func (p *THttpClient) Write(buf []byte) (int, error) {
	n, err := p.requestBuffer.Write(buf)
	return n, err
}

func (p *THttpClient) WriteByte(c byte) error {
	return p.requestBuffer.WriteByte(c)
}

func (p *THttpClient) WriteString(s string) (n int, err error) {
	return p.requestBuffer.WriteString(s)
}

func (p *THttpClient) Flush() error {
	client := &http.Client{}
	req, err := http.NewRequest("POST", p.url.String(), p.requestBuffer)
	if err != nil {
		return NewTTransportExceptionFromError(err)
	}
	p.header.Add("Content-Type", "application/x-thrift")
	req.Header = p.header
	response, err := client.Do(req)
	if err != nil {
		return NewTTransportExceptionFromError(err)
	}
	if response.StatusCode != http.StatusOK {
		// TODO(pomack) log bad response
		return NewTTransportException(UNKNOWN_TRANSPORT_EXCEPTION, "HTTP Response code: "+strconv.Itoa(response.StatusCode))
	}
	p.response = response
	return nil
}
<commit_msg>lib/go: Don't write multiple time the header in the http client transport<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package thrift

import (
	"bytes"
	"io"
	"net/http"
	"net/url"
	"strconv"
)

type THttpClient struct {
	response           *http.Response
	url                *url.URL
	requestBuffer      *bytes.Buffer
	header             http.Header
	nsecConnectTimeout int64
	nsecReadTimeout    int64
}

type THttpClientTransportFactory struct {
	url    string
	isPost bool
}

func (p *THttpClientTransportFactory) GetTransport(trans TTransport) TTransport {
	if trans != nil {
		t, ok := trans.(*THttpClient)
		if ok && t.url != nil {
			if t.requestBuffer != nil {
				t2, _ := NewTHttpPostClient(t.url.String())
				return t2
			}
			t2, _ := NewTHttpClient(t.url.String())
			return t2
		}
	}
	if p.isPost {
		s, _ := NewTHttpPostClient(p.url)
		return s
	}
	s, _ := NewTHttpClient(p.url)
	return s
}

func NewTHttpClientTransportFactory(url string) *THttpClientTransportFactory {
	return &THttpClientTransportFactory{url: url, isPost: false}
}

func NewTHttpPostClientTransportFactory(url string) *THttpClientTransportFactory {
	return &THttpClientTransportFactory{url: url, isPost: true}
}

func NewTHttpClient(urlstr string) (TTransport, error) {
	parsedURL, err := url.Parse(urlstr)
	if err != nil {
		return nil, err
	}
	response, err := http.Get(urlstr)
	if err != nil {
		return nil, err
	}
	return &THttpClient{response: response, url: parsedURL}, nil
}

func NewTHttpPostClient(urlstr string) (TTransport, error) {
	parsedURL, err := url.Parse(urlstr)
	if err != nil {
		return nil, err
	}
	buf := make([]byte, 0, 1024)
	return &THttpClient{
		url:           parsedURL,
		requestBuffer: bytes.NewBuffer(buf),
		header: http.Header{
			"Content-Type": []string{"application/x-thrift"},
		},
	}, nil
}

// Set the HTTP Header for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// httpTrans.SetHeader("User-Agent","Thrift Client 1.0")
func (p *THttpClient) SetHeader(key string, value string) {
	p.header.Add(key, value)
}

// Get the HTTP Header represented by the supplied Header Key for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// hdrValue := httpTrans.GetHeader("User-Agent")
func (p *THttpClient) GetHeader(key string) string {
	return p.header.Get(key)
}

// Deletes the HTTP Header given a Header Key for this specific Thrift Transport
// It is important that you first assert the TTransport as a THttpClient type
// like so:
//
// httpTrans := trans.(THttpClient)
// httpTrans.DelHeader("User-Agent")
func (p *THttpClient) DelHeader(key string) {
	p.header.Del(key)
}

func (p *THttpClient) Open() error {
	// do nothing
	return nil
}

func (p *THttpClient) IsOpen() bool {
	return p.response != nil || p.requestBuffer != nil
}

func (p *THttpClient) Close() error {
	if p.response != nil && p.response.Body != nil {
		err := p.response.Body.Close()
		p.response = nil
		return err
	}
	if p.requestBuffer != nil {
		p.requestBuffer.Reset()
		p.requestBuffer = nil
	}
	return nil
}

func (p *THttpClient) Read(buf []byte) (int, error) {
	if p.response == nil {
		return 0, NewTTransportException(NOT_OPEN, "Response buffer is empty, no request.")
	}
	n, err := p.response.Body.Read(buf)
	if n > 0 && (err == nil || err == io.EOF) {
		return n, nil
	}
	return n, NewTTransportExceptionFromError(err)
}

func (p *THttpClient) ReadByte() (c byte, err error) {
	return readByte(p.response.Body)
}

func (p *THttpClient) Write(buf []byte) (int, error) {
	n, err := p.requestBuffer.Write(buf)
	return n, err
}

func (p *THttpClient) WriteByte(c byte) error {
	return p.requestBuffer.WriteByte(c)
}

func (p *THttpClient) WriteString(s string) (n int, err error) {
	return p.requestBuffer.WriteString(s)
}

func (p *THttpClient) Flush() error {
	client := &http.Client{}
	req, err := http.NewRequest("POST", p.url.String(), p.requestBuffer)
	if err != nil {
		return NewTTransportExceptionFromError(err)
	}
	req.Header = p.header
	response, err := client.Do(req)
	if err != nil {
		return NewTTransportExceptionFromError(err)
	}
	if response.StatusCode != http.StatusOK {
		// TODO(pomack) log bad response
		return NewTTransportException(UNKNOWN_TRANSPORT_EXCEPTION, "HTTP Response code: "+strconv.Itoa(response.StatusCode))
	}
	p.response = response
	return nil
}
<|endoftext|>


<commit_before>package task

import "testing"

func TestTaskState(t *testing.T) {
	dummyTask := NewTask("dummy", "foo", "bar")
	got := dummyTask.State
	want := TaskStateUnknown
	if want != got {
		t.Errorf("Incorrect task state: want %q, got %q", want, got)
	}
}

func TestTaskStateCommand(t *testing.T) {
	dummyTask := NewTask("dummy", "foo", "bar")
	got := dummyTask.Command
	want := "dummy"
	if want != got {
		t.Errorf("Incorrect task command: want %q, got %q", want, got)
	}
}

<commit_msg>Fix unit test function name<commit_after>package task

import "testing"

func TestTaskState(t *testing.T) {
	dummyTask := NewTask("dummy", "foo", "bar")
	got := dummyTask.State
	want := TaskStateUnknown
	if want != got {
		t.Errorf("Incorrect task state: want %q, got %q", want, got)
	}
}

func TestTaskCommand(t *testing.T) {
	dummyTask := NewTask("dummy", "foo", "bar")
	got := dummyTask.Command
	want := "dummy"
	if want != got {
		t.Errorf("Incorrect task command: want %q, got %q", want, got)
	}
}

<|endoftext|>


<commit_before>package template

import (
	"bytes"
	"fmt"
	"gnd.la/log"
	"gnd.la/util/hashutil"
	"io/ioutil"
	"path"
	"text/template"
)

func executeAsset(t *Template, name string) (string, error) {
	log.Debugf("executing asset template %s", name)
	f, _, err := t.AssetsManager.Load(name)
	if err != nil {
		return "", err
	}
	defer f.Close()
	data, err := ioutil.ReadAll(f)
	if err != nil {
		return "", err
	}
	tmpl := template.New(name)
	if _, err := tmpl.Parse(string(data)); err != nil {
		return "", err
	}
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, t.vars); err != nil {
		return "", err
	}
	ext := path.Ext(name)
	nonExt := name[:len(name)-len(ext)]
	f.Seek(0, 0)
	out := fmt.Sprintf("%s.%s%s", nonExt, hashutil.Fnv32a(f), ext)
	w, err := t.AssetsManager.Create(out, true)
	if err != nil {
		return "", err
	}
	if _, err := w.Write(buf.Bytes()); err != nil {
		return "", err
	}
	if err := w.Close(); err != nil {
		return "", err
	}
	return out, nil
}
<commit_msg>Pass functions to templated assets<commit_after>package template

import (
	"bytes"
	"fmt"
	"gnd.la/log"
	"gnd.la/util/hashutil"
	"io/ioutil"
	"path"
	"text/template"
)

func executeAsset(t *Template, name string) (string, error) {
	log.Debugf("executing asset template %s", name)
	f, _, err := t.AssetsManager.Load(name)
	if err != nil {
		return "", err
	}
	defer f.Close()
	data, err := ioutil.ReadAll(f)
	if err != nil {
		return "", err
	}
	tmpl := template.New(name)
	tmpl.Funcs(template.FuncMap(templateFuncs))
	tmpl.Funcs(template.FuncMap(t.funcMap))
	if _, err := tmpl.Parse(string(data)); err != nil {
		return "", err
	}
	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, t.vars); err != nil {
		return "", err
	}
	ext := path.Ext(name)
	nonExt := name[:len(name)-len(ext)]
	f.Seek(0, 0)
	out := fmt.Sprintf("%s.%s%s", nonExt, hashutil.Fnv32a(f), ext)
	w, err := t.AssetsManager.Create(out, true)
	if err != nil {
		return "", err
	}
	if _, err := w.Write(buf.Bytes()); err != nil {
		return "", err
	}
	if err := w.Close(); err != nil {
		return "", err
	}
	return out, nil
}
<|endoftext|>