<commit_before>package daemon

import (
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"syscall"
)

// Service Daemon
type ServiceDaemon struct {
	Daemon
}

func (daemon *ServiceDaemon) Manage(service Service) (string, error) {

	usage := "Usage: myservice install | remove | start | stop | status"

	// if received any kind of command, do it
	if len(os.Args) > 1 {
		command := os.Args[1]
		switch command {
		case "install":
			return daemon.Install()
		case "remove":
			return daemon.Remove()
		case "start":
			return daemon.Start()
		case "stop":
			return daemon.Stop()
		case "status":
			return daemon.Status()
		default:
			return usage, nil
		}
	}

	process := service.GetProcess()
	process()

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

	// Set up listener for defined host and port
	listener, err := net.Listen("tcp", ":"+strconv.Itoa(service.GetPort()))
	if err != nil {
		return "Possibly was a problem with the port binding", err
	}

	// set up channel on which to send accepted connections
	listen := make(chan net.Conn, 100)
	go acceptConnection(listener, listen)

	// loop work cycle with accept connections or interrupt
	// by system signal
	for {
		select {
		case conn := <-listen:
			go handleClient(conn)
		case killSignal := <-interrupt:
			log.Println("Got signal:", killSignal)
			log.Println("Stoping listening on ", listener.Addr())
			listener.Close()
			if killSignal == os.Interrupt {
				return "Daemon was interruped by system signal", nil
			}
			return "Daemon was killed", nil
		}
	}

	// never happen, but need to complete code
	return usage, nil
}

// Accept a client connection and collect it in a channel
func acceptConnection(listener net.Listener, listen chan<- net.Conn) {
	for {
		conn, err := listener.Accept()
		if err != nil {
			continue
		}
		listen <- conn
	}
}

func handleClient(client net.Conn) {
	for {
		buf := make([]byte, 4096)
		numbytes, err := client.Read(buf)
		if numbytes == 0 || err != nil {
			return
		}
		client.Write(buf)
	}
}

// Services

type Service struct {
	Name        string
	Description string
	Port        int
	Process     func()
}

func (service Service) GetName() string {
	return service.Name
}

func (service Service) GetPort() int {
	return service.Port
}

func (service Service) GetDescription() string {
	return service.Description
}

func (service Service) GetProcess() func() {
	return service.Process
}

func (service Service) Daemon() {

	srv, err := New(service.GetName(), service.GetDescription())

	if err != nil {
		fmt.Println("Error: ", err)
		os.Exit(1)
	}

	service_daemon := &ServiceDaemon{srv}
	status, err := service_daemon.Manage(service)

	if err != nil {
		fmt.Println(status, "\nError: ", err)
		os.Exit(1)
	}

	fmt.Println(status)
}
<commit_msg>Update initialize function<commit_after>package daemon

import (
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"syscall"
)

// Service Daemon
type ServiceDaemon struct {
	Daemon
}

func (daemon *ServiceDaemon) Manage(service Service) (string, error) {

	usage := "Usage: myservice install | remove | start | stop | status"

	// if received any kind of command, do it
	if len(os.Args) > 1 {
		command := os.Args[1]
		switch command {
		case "install":
			return daemon.Install()
		case "remove":
			return daemon.Remove()
		case "start":
			return daemon.Start()
		case "stop":
			return daemon.Stop()
		case "status":
			return daemon.Status()
		default:
			return usage, nil
		}
	}

	process := service.GetProcess()
	process()

	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt, os.Kill, syscall.SIGTERM)

	// Set up listener for defined host and port
	listener, err := net.Listen("tcp", ":"+strconv.Itoa(service.GetPort()))
	if err != nil {
		return "Possibly was a problem with the port binding", err
	}

	// set up channel on which to send accepted connections
	listen := make(chan net.Conn, 100)
	go acceptConnection(listener, listen)

	// loop work cycle with accept connections or interrupt
	// by system signal
	for {
		select {
		case conn := <-listen:
			go handleClient(conn)
		case killSignal := <-interrupt:
			log.Println("Got signal:", killSignal)
			log.Println("Stoping listening on ", listener.Addr())
			listener.Close()
			if killSignal == os.Interrupt {
				return "Daemon was interruped by system signal", nil
			}
			return "Daemon was killed", nil
		}
	}

	// never happen, but need to complete code
	return usage, nil
}

// Accept a client connection and collect it in a channel
func acceptConnection(listener net.Listener, listen chan<- net.Conn) {
	for {
		conn, err := listener.Accept()
		if err != nil {
			continue
		}
		listen <- conn
	}
}

func handleClient(client net.Conn) {
	for {
		buf := make([]byte, 4096)
		numbytes, err := client.Read(buf)
		if numbytes == 0 || err != nil {
			return
		}
		client.Write(buf)
	}
}

// Services

type Service struct {
	Name        string
	Description string
	Port        int
	Process     func()
}

func (service Service) GetName() string {
	return service.Name
}

func (service Service) GetPort() int {
	return service.Port
}

func (service Service) GetDescription() string {
	return service.Description
}

func (service Service) GetProcess() func() {
	return service.Process
}

func (service Service) Initialize() {

	srv, err := New(service.GetName(), service.GetDescription())

	if err != nil {
		fmt.Println("Error: ", err)
		os.Exit(1)
	}

	service_daemon := &ServiceDaemon{srv}
	status, err := service_daemon.Manage(service)

	if err != nil {
		fmt.Println(status, "\nError: ", err)
		os.Exit(1)
	}

	fmt.Println(status)
}
<|endoftext|>


<commit_before>package core

import (
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/mdlayher/wavepipe/data"

	"github.com/go-martini/martini"
	"github.com/martini-contrib/render"
)

// Martini instance to test against
var m = martini.New()

func init() {
	// Set up database connection
	data.DB = new(data.SqliteBackend)
	data.DB.DSN("~/.config/wavepipe/wavepipe.db")
	if err := data.DB.Open(); err != nil {
		os.Exit(1)
	}

	// Map a dummy user for some API calls
	m.Use(render.Renderer(render.Options{}))
	m.Use(func(c martini.Context) {
		c.Map(new(data.User))
	})

	// Set up Martini with API routes
	r := martini.NewRouter()
	r.Group("/api/:version", apiRoutes)
	m.Action(r.Handle)

}

// TestAPIRouter verifies that all API request processing functionality is working properly
func TestAPIRouter(t *testing.T) {
	// Table of tests to run, and their expected HTTP status results
	var tests = []struct {
		code int
		url  string
	}{
		// Albums API
		//   - valid request
		{200, "/api/v0/albums"},
		//   - valid request for 1 item
		{200, "/api/v0/albums/1"},
		//   - valid limit items request
		{200, "/api/v0/albums?limit=0,10"},
		//   - invalid API version
		{400, "/api/v999/albums"},
		//   - invalid integer album ID
		{400, "/api/v0/albums/foo"},
		//   - missing second integer for limit
		{400, "/api/v0/albums?limit=0"},
		//   - invalid integer pair for limit
		{400, "/api/v0/albums?limit=foo,bar"},
		//   - album ID not found
		{404, "/api/v0/albums/99999999"},

		// Art API - skip valid requests, due to binary output
		//   - invalid API version
		{400, "/api/v999/art"},
		//   - no integer ID provided
		{400, "/api/v0/art"},
		//   - invalid art stream ID
		{400, "/api/v0/art/foo"},
		//   - art ID not found
		{404, "/api/v0/art/99999999"},

		// Artists API
		//   - valid request
		{200, "/api/v0/artists"},
		//   - valid request for 1 item
		{200, "/api/v0/artists/1"},
		//   - valid limit items request
		{200, "/api/v0/artists?limit=0,10"},
		//   - invalid API version
		{400, "/api/v999/artists"},
		//   - invalid integer artist ID
		{400, "/api/v0/artists/foo"},
		//   - missing second integer for limit
		{400, "/api/v0/artists?limit=0"},
		//   - invalid integer pair for limit
		{400, "/api/v0/artists?limit=foo,bar"},
		//   - artist ID not found
		{404, "/api/v0/artists/99999999"},

		// Folders API
		//   - valid request
		{200, "/api/v0/folders"},
		//   - valid request for 1 item
		{200, "/api/v0/folders/1"},
		//   - invalid API version
		{400, "/api/v999/folders"},
		//   - valid limit items request
		{200, "/api/v0/folders?limit=0,10"},
		//   - invalid integer folder ID
		{400, "/api/v0/folders/foo"},
		//   - missing second integer for limit
		{400, "/api/v0/folders?limit=0"},
		//   - invalid integer pair for limit
		{400, "/api/v0/folders?limit=foo,bar"},
		//   - folder ID not found
		{404, "/api/v0/folders/99999999"},

		// LastFM API - skip valid requests, due to need for external service
		//   - invalid API version
		{400, "/api/v999/lastfm"},
		//   - no string action provided
		{400, "/api/v0/lastfm"},
		//   - invalid string action provided
		{400, "/api/v0/lastfm/foo"},
		//   - login: no username provided
		{400, "/api/v0/lastfm/login"},
		//   - login: no password provided
		{400, "/api/v0/lastfm/login?lfmu=foo"},
		//   - action: user must authenticate to last.fm
		{401, "/api/v0/lastfm/nowplaying"},
		//   - action: user must authenticate to last.fm
		{401, "/api/v0/lastfm/scrobble"},
		// Cannot test other calls without a valid Last.fm token

		// Login/Logout API - skip due to need for sessions and users

		// Search API
		//   - valid request
		{200, "/api/v0/search/foo"},
		//   - invalid API version
		{400, "/api/v999/search"},
		//   - no search query specified
		{400, "/api/v0/search"},

		// Songs API
		//   - valid request
		{200, "/api/v0/songs"},
		//   - valid request for 1 item
		{200, "/api/v0/songs/1"},
		//   - valid limit items request
		{200, "/api/v0/songs?limit=0,10"},
		//   - valid random items request
		{200, "/api/v0/songs?random=10"},
		//   - invalid API version
		{400, "/api/v999/songs"},
		//   - invalid integer song ID
		{400, "/api/v0/songs/foo"},
		//   - missing second integer for limit
		{400, "/api/v0/songs?limit=0"},
		//   - invalid integer pair for limit
		{400, "/api/v0/songs?limit=foo,bar"},
		//   - invalid integer for random
		{400, "/api/v0/songs?random=foo"},
		//   - song ID not found
		{404, "/api/v0/songs/99999999"},

		// Status API
		//   - valid request
		{200, "/api/v0/status"},
		//   - invalid API version
		{400, "/api/v999/status"},

		// Stream API - skip valid requests, due to binary output
		//   - invalid API version
		{400, "/api/v999/stream"},
		//   - no integer stream ID provided
		{400, "/api/v0/stream"},
		//   - invalid stream stream ID
		{400, "/api/v0/stream/foo"},
		//   - song ID not found
		{404, "/api/v0/stream/99999999"},

		// Transcode API - skip valid requests, due to binary output
		//   - invalid API version
		{400, "/api/v999/transcode"},
		//   - no integer transcode ID provided
		{400, "/api/v0/transcode"},
		//   - invalid transcode transcode ID
		{400, "/api/v0/transcode/foo"},
		//   - song ID not found
		{404, "/api/v0/transcode/99999999"},
		//   - ffmpeg not found, transcoding disabled
		{503, "/api/v0/transcode/1"},
	}

	// Iterate all tests
	for _, test := range tests {
		// Generate a new HTTP request
		r, err := http.NewRequest("GET", "http://localhost:8080"+test.url, nil)
		if err != nil {
			t.Fatalf("Failed to create HTTP request")
		}

		// Capture HTTP response via recorder
		w := httptest.NewRecorder()

		// Perform request
		m.ServeHTTP(w, r)

		// Validate results
		if w.Code != test.code {
			t.Fatalf("[%v != %v] %s", w.Code, test.code, test.url)
		}

		log.Printf("OK: [%d] %s", test.code, test.url)
	}
}
<commit_msg>core/apiRouter_test: add api/status?metrics=true test<commit_after>package core

import (
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/mdlayher/wavepipe/data"

	"github.com/go-martini/martini"
	"github.com/martini-contrib/render"
)

// Martini instance to test against
var m = martini.New()

func init() {
	// Set up database connection
	data.DB = new(data.SqliteBackend)
	data.DB.DSN("~/.config/wavepipe/wavepipe.db")
	if err := data.DB.Open(); err != nil {
		os.Exit(1)
	}

	// Map a dummy user for some API calls
	m.Use(render.Renderer(render.Options{}))
	m.Use(func(c martini.Context) {
		c.Map(new(data.User))
	})

	// Set up Martini with API routes
	r := martini.NewRouter()
	r.Group("/api/:version", apiRoutes)
	m.Action(r.Handle)

}

// TestAPIRouter verifies that all API request processing functionality is working properly
func TestAPIRouter(t *testing.T) {
	// Table of tests to run, and their expected HTTP status results
	var tests = []struct {
		code int
		url  string
	}{
		// Albums API
		//   - valid request
		{200, "/api/v0/albums"},
		//   - valid request for 1 item
		{200, "/api/v0/albums/1"},
		//   - valid limit items request
		{200, "/api/v0/albums?limit=0,10"},
		//   - invalid API version
		{400, "/api/v999/albums"},
		//   - invalid integer album ID
		{400, "/api/v0/albums/foo"},
		//   - missing second integer for limit
		{400, "/api/v0/albums?limit=0"},
		//   - invalid integer pair for limit
		{400, "/api/v0/albums?limit=foo,bar"},
		//   - album ID not found
		{404, "/api/v0/albums/99999999"},

		// Art API - skip valid requests, due to binary output
		//   - invalid API version
		{400, "/api/v999/art"},
		//   - no integer ID provided
		{400, "/api/v0/art"},
		//   - invalid art stream ID
		{400, "/api/v0/art/foo"},
		//   - art ID not found
		{404, "/api/v0/art/99999999"},

		// Artists API
		//   - valid request
		{200, "/api/v0/artists"},
		//   - valid request for 1 item
		{200, "/api/v0/artists/1"},
		//   - valid limit items request
		{200, "/api/v0/artists?limit=0,10"},
		//   - invalid API version
		{400, "/api/v999/artists"},
		//   - invalid integer artist ID
		{400, "/api/v0/artists/foo"},
		//   - missing second integer for limit
		{400, "/api/v0/artists?limit=0"},
		//   - invalid integer pair for limit
		{400, "/api/v0/artists?limit=foo,bar"},
		//   - artist ID not found
		{404, "/api/v0/artists/99999999"},

		// Folders API
		//   - valid request
		{200, "/api/v0/folders"},
		//   - valid request for 1 item
		{200, "/api/v0/folders/1"},
		//   - invalid API version
		{400, "/api/v999/folders"},
		//   - valid limit items request
		{200, "/api/v0/folders?limit=0,10"},
		//   - invalid integer folder ID
		{400, "/api/v0/folders/foo"},
		//   - missing second integer for limit
		{400, "/api/v0/folders?limit=0"},
		//   - invalid integer pair for limit
		{400, "/api/v0/folders?limit=foo,bar"},
		//   - folder ID not found
		{404, "/api/v0/folders/99999999"},

		// LastFM API - skip valid requests, due to need for external service
		//   - invalid API version
		{400, "/api/v999/lastfm"},
		//   - no string action provided
		{400, "/api/v0/lastfm"},
		//   - invalid string action provided
		{400, "/api/v0/lastfm/foo"},
		//   - login: no username provided
		{400, "/api/v0/lastfm/login"},
		//   - login: no password provided
		{400, "/api/v0/lastfm/login?lfmu=foo"},
		//   - action: user must authenticate to last.fm
		{401, "/api/v0/lastfm/nowplaying"},
		//   - action: user must authenticate to last.fm
		{401, "/api/v0/lastfm/scrobble"},
		// Cannot test other calls without a valid Last.fm token

		// Login/Logout API - skip due to need for sessions and users

		// Search API
		//   - valid request
		{200, "/api/v0/search/foo"},
		//   - invalid API version
		{400, "/api/v999/search"},
		//   - no search query specified
		{400, "/api/v0/search"},

		// Songs API
		//   - valid request
		{200, "/api/v0/songs"},
		//   - valid request for 1 item
		{200, "/api/v0/songs/1"},
		//   - valid limit items request
		{200, "/api/v0/songs?limit=0,10"},
		//   - valid random items request
		{200, "/api/v0/songs?random=10"},
		//   - invalid API version
		{400, "/api/v999/songs"},
		//   - invalid integer song ID
		{400, "/api/v0/songs/foo"},
		//   - missing second integer for limit
		{400, "/api/v0/songs?limit=0"},
		//   - invalid integer pair for limit
		{400, "/api/v0/songs?limit=foo,bar"},
		//   - invalid integer for random
		{400, "/api/v0/songs?random=foo"},
		//   - song ID not found
		{404, "/api/v0/songs/99999999"},

		// Status API
		//   - valid request
		{200, "/api/v0/status"},
		//   - valid request with metrics
		{200, "/api/v0/status?metrics=true"},
		//   - invalid API version
		{400, "/api/v999/status"},

		// Stream API - skip valid requests, due to binary output
		//   - invalid API version
		{400, "/api/v999/stream"},
		//   - no integer stream ID provided
		{400, "/api/v0/stream"},
		//   - invalid stream stream ID
		{400, "/api/v0/stream/foo"},
		//   - song ID not found
		{404, "/api/v0/stream/99999999"},

		// Transcode API - skip valid requests, due to binary output
		//   - invalid API version
		{400, "/api/v999/transcode"},
		//   - no integer transcode ID provided
		{400, "/api/v0/transcode"},
		//   - invalid transcode transcode ID
		{400, "/api/v0/transcode/foo"},
		//   - song ID not found
		{404, "/api/v0/transcode/99999999"},
		//   - ffmpeg not found, transcoding disabled
		{503, "/api/v0/transcode/1"},
	}

	// Iterate all tests
	for _, test := range tests {
		// Generate a new HTTP request
		r, err := http.NewRequest("GET", "http://localhost:8080"+test.url, nil)
		if err != nil {
			t.Fatalf("Failed to create HTTP request")
		}

		// Capture HTTP response via recorder
		w := httptest.NewRecorder()

		// Perform request
		m.ServeHTTP(w, r)

		// Validate results
		if w.Code != test.code {
			t.Fatalf("[%v != %v] %s", w.Code, test.code, test.url)
		}

		log.Printf("OK: [%d] %s", test.code, test.url)
	}
}
<|endoftext|>


<commit_before>package instruments

import (
	"encoding/json"
	"errors"
	"net/http"

	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry-incubator/metricz/instrumentation"
	"github.com/cloudfoundry/gunk/urljoiner"
	"github.com/pivotal-golang/lager"
)

type Leader struct {
	statsEndpoint string

	logger lager.Logger
}

var ErrRedirected = errors.New("redirected to leader")

func NewLeader(etcdAddr string, logger lager.Logger) *Leader {
	return &Leader{
		statsEndpoint: urljoiner.Join(etcdAddr, "v2", "stats", "leader"),

		logger: logger,
	}
}

func (leader *Leader) Emit() instrumentation.Context {
	context := instrumentation.Context{
		Name:    "leader",
		Metrics: []instrumentation.Metric{},
	}

	var stats RaftFollowersStats

	client := cf_http.NewClient()
	client.CheckRedirect = func(*http.Request, []*http.Request) error {
		return ErrRedirected
	}

	resp, err := client.Get(leader.statsEndpoint)
	if err != nil {
		leader.logger.Error("failed-to-collect-leader-stats", err)
		return context
	}

	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(&stats)
	if err != nil {
		leader.logger.Error("failed-to-unmarshal-leader-stats", err)
		return context
	}

	context.Metrics = []instrumentation.Metric{
		{
			Name:  "Followers",
			Value: len(stats.Followers),
		},
	}

	for name, follower := range stats.Followers {
		context.Metrics = append(context.Metrics, instrumentation.Metric{
			Name:  "Latency",
			Value: follower.Latency.Current,
			Tags: map[string]interface{}{
				"follower": name,
			},
		})
	}

	return context
}
<commit_msg>Create a single http client<commit_after>package instruments

import (
	"encoding/json"
	"errors"
	"net/http"

	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry-incubator/metricz/instrumentation"
	"github.com/cloudfoundry/gunk/urljoiner"
	"github.com/pivotal-golang/lager"
)

type Leader struct {
	statsEndpoint string
	client        *http.Client

	logger lager.Logger
}

var ErrRedirected = errors.New("redirected to leader")

func NewLeader(etcdAddr string, logger lager.Logger) *Leader {
	client := cf_http.NewClient()
	client.CheckRedirect = func(*http.Request, []*http.Request) error {
		return ErrRedirected
	}

	return &Leader{
		statsEndpoint: urljoiner.Join(etcdAddr, "v2", "stats", "leader"),
		client:        client,

		logger: logger,
	}
}

func (leader *Leader) Emit() instrumentation.Context {
	context := instrumentation.Context{
		Name:    "leader",
		Metrics: []instrumentation.Metric{},
	}

	var stats RaftFollowersStats

	resp, err := leader.client.Get(leader.statsEndpoint)
	if err != nil {
		leader.logger.Error("failed-to-collect-leader-stats", err)
		return context
	}

	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(&stats)
	if err != nil {
		leader.logger.Error("failed-to-unmarshal-leader-stats", err)
		return context
	}

	context.Metrics = []instrumentation.Metric{
		{
			Name:  "Followers",
			Value: len(stats.Followers),
		},
	}

	for name, follower := range stats.Followers {
		context.Metrics = append(context.Metrics, instrumentation.Metric{
			Name:  "Latency",
			Value: follower.Latency.Current,
			Tags: map[string]interface{}{
				"follower": name,
			},
		})
	}

	return context
}
<|endoftext|>


<commit_before>// package pstree provides an API to retrieve the process tree of a given
// process-id.
package pstree

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// New returns the whole system process tree.
func New() (*Tree, error) {
	files, err := filepath.Glob("/proc/[0-9]*")
	if err != nil {
		return nil, err
	}

	procs := make(map[int]Process, len(files))
	for _, dir := range files {
		proc, err := scan(dir)
		if err != nil {
			return nil, err
		}
		if proc.Pid == 0 {
			// process vanished since Glob.
			continue
		}
		procs[proc.Pid] = proc
	}

	for pid, proc := range procs {
		if proc.Parent == 0 {
			continue
		}
		parent, ok := procs[proc.Parent]
		if !ok {
			log.Panicf(
				"internal logic error. parent of [%d] does not exist!",
				pid,
			)
		}
		parent.Children = append(parent.Children, pid)
		procs[parent.Pid] = parent
	}

	for pid, proc := range procs {
		if len(proc.Children) > 0 {
			sort.Ints(proc.Children)
		}
		procs[pid] = proc
	}

	tree := &Tree{
		Procs: procs,
	}
	return tree, err
}

const (
	statfmt = "%d %s %c %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d"
)

func scan(dir string) (Process, error) {
	f, err := os.Open(filepath.Join(dir, "stat"))
	if err != nil {
		// process vanished since Glob.
		return Process{}, nil
	}
	defer f.Close()

	// see: http://man7.org/linux/man-pages/man5/proc.5.html
	stat := struct {
		pid       int    // process ID
		comm      string // filename of the executable in parentheses
		state     byte   // process state
		ppid      int    // pid of the parent process
		pgrp      int    // process group ID of the process
		session   int    // session ID of the process
		tty       int    // controlling terminal of the process
		tpgid     int    // ID of foreground process group
		flags     uint32 // kernel flags word of the process
		minflt    uint64 // number of minor faults the process has made which have not required loading a memory page from disk
		cminflt   uint64 // number of minor faults the process's waited-for children have made
		majflt    uint64 // number of major faults the process has made which have required loading a memory page from disk
		cmajflt   uint64 // number of major faults the process's waited-for children have made
		utime     uint64 // user time in clock ticks
		stime     uint64 // system time in clock ticks
		cutime    int64  // children user time in clock ticks
		cstime    int64  // children system time in clock ticks
		priority  int64  // priority
		nice      int64  // the nice value
		nthreads  int64  // number of threads in this process
		itrealval int64  // time in jiffies before next SIGALRM is sent to the process dure to an interval timer
		starttime int64  // time the process started after system boot in clock ticks
		vsize     uint64 // virtual memory size in bytes
		rss       int64  // resident set size: number of pages the process has in real memory
	}{}

	_, err = fmt.Fscanf(
		f, statfmt,
		&stat.pid, &stat.comm, &stat.state,
		&stat.ppid, &stat.pgrp, &stat.session,
		&stat.tty, &stat.tpgid, &stat.flags,
		&stat.minflt, &stat.cminflt, &stat.majflt, &stat.cmajflt,
		&stat.utime, &stat.stime,
		&stat.cutime, &stat.cstime,
		&stat.priority,
		&stat.nice,
		&stat.nthreads,
		&stat.itrealval, &stat.starttime,
		&stat.vsize, &stat.rss,
	)
	if err != nil {
		return Process{}, err
	}

	name := stat.comm
	if strings.HasPrefix(name, "(") && strings.HasSuffix(name, ")") {
		name = name[1 : len(name)-1]
	}
	return Process{
		Name:   name,
		Pid:    stat.pid,
		Parent: stat.ppid,
	}, err
}

// Tree is a tree of processes.
type Tree struct {
	Procs map[int]Process
}

// Process stores informations about a UNIX process
type Process struct {
	Name     string
	Pid      int
	Parent   int
	Children []int
}
<commit_msg>pstree: expose stat<commit_after>// package pstree provides an API to retrieve the process tree from procfs.
package pstree

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

// New returns the whole system process tree.
func New() (*Tree, error) {
	files, err := filepath.Glob("/proc/[0-9]*")
	if err != nil {
		return nil, err
	}

	procs := make(map[int]Process, len(files))
	for _, dir := range files {
		proc, err := scan(dir)
		if err != nil {
			return nil, err
		}
		if proc.Stat.Pid == 0 {
			// process vanished since Glob.
			continue
		}
		procs[proc.Stat.Pid] = proc
	}

	for pid, proc := range procs {
		if proc.Stat.Ppid == 0 {
			continue
		}
		parent, ok := procs[proc.Stat.Ppid]
		if !ok {
			log.Panicf(
				"internal logic error. parent of [%d] does not exist!",
				pid,
			)
		}
		parent.Children = append(parent.Children, pid)
		procs[parent.Stat.Pid] = parent
	}

	for pid, proc := range procs {
		if len(proc.Children) > 0 {
			sort.Ints(proc.Children)
		}
		procs[pid] = proc
	}

	tree := &Tree{
		Procs: procs,
	}
	return tree, err
}

const (
	statfmt = "%d %s %c %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d %d"
)

// ProcessStat contains process information.
// see: http://man7.org/linux/man-pages/man5/proc.5.html
type ProcessStat struct {
	Pid       int    // process ID
	Comm      string // filename of the executable in parentheses
	State     byte   // process state
	Ppid      int    // pid of the parent process
	Pgrp      int    // process group ID of the process
	Session   int    // session ID of the process
	Tty       int    // controlling terminal of the process
	Tpgid     int    // ID of foreground process group
	Flags     uint32 // kernel flags word of the process
	Minflt    uint64 // number of minor faults the process has made which have not required loading a memory page from disk
	Cminflt   uint64 // number of minor faults the process's waited-for children have made
	Majflt    uint64 // number of major faults the process has made which have required loading a memory page from disk
	Cmajflt   uint64 // number of major faults the process's waited-for children have made
	Utime     uint64 // user time in clock ticks
	Stime     uint64 // system time in clock ticks
	Cutime    int64  // children user time in clock ticks
	Cstime    int64  // children system time in clock ticks
	Priority  int64  // priority
	Nice      int64  // the nice value
	Nthreads  int64  // number of threads in this process
	Itrealval int64  // time in jiffies before next SIGALRM is sent to the process due to an interval timer
	Starttime int64  // time the process started after system boot in clock ticks
	Vsize     uint64 // virtual memory size in bytes
	Rss       int64  // resident set size: number of pages the process has in real memory
}

func scan(dir string) (Process, error) {
	f, err := os.Open(filepath.Join(dir, "stat"))
	if err != nil {
		// process vanished since Glob.
		return Process{}, nil
	}
	defer f.Close()

	var proc Process
	_, err = fmt.Fscanf(
		f, statfmt,
		&proc.Stat.Pid, &proc.Stat.Comm, &proc.Stat.State,
		&proc.Stat.Ppid, &proc.Stat.Pgrp, &proc.Stat.Session,
		&proc.Stat.Tty, &proc.Stat.Tpgid, &proc.Stat.Flags,
		&proc.Stat.Minflt, &proc.Stat.Cminflt, &proc.Stat.Majflt, &proc.Stat.Cmajflt,
		&proc.Stat.Utime, &proc.Stat.Stime,
		&proc.Stat.Cutime, &proc.Stat.Cstime,
		&proc.Stat.Priority,
		&proc.Stat.Nice,
		&proc.Stat.Nthreads,
		&proc.Stat.Itrealval, &proc.Stat.Starttime,
		&proc.Stat.Vsize, &proc.Stat.Rss,
	)
	if err != nil {
		return proc, err
	}

	proc.Name = proc.Stat.Comm
	if strings.HasPrefix(proc.Name, "(") && strings.HasSuffix(proc.Name, ")") {
		proc.Name = proc.Name[1 : len(proc.Name)-1]
	}
	return proc, nil
}

// Tree is a tree of processes.
type Tree struct {
	Procs map[int]Process
}

// Process stores information about a UNIX process.
type Process struct {
	Name     string
	Stat     ProcessStat
	Children []int
}
<|endoftext|>


<commit_before>package commands

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	"github.com/ipfs/go-ipfs/blocks"
	util "github.com/ipfs/go-ipfs/blocks/blockstore/util"
	cmds "github.com/ipfs/go-ipfs/commands"

	cid "gx/ipfs/QmV5gPoRsjN1Gid3LMdNZTyfCtP2DsvqEbMAmz82RmmiGk/go-cid"
	u "gx/ipfs/QmZuY8aV7zbNXVy6DyN9SmnuH3o9nG852F4aTiSBpts8d1/go-ipfs-util"
	mh "gx/ipfs/QmbZ6Cee2uHjG7hf19qLHppgKDRtaG4CVtMzdmK9VCVqLu/go-multihash"
)

type BlockStat struct {
	Key  string
	Size int
}

func (bs BlockStat) String() string {
	return fmt.Sprintf("Key: %s\nSize: %d\n", bs.Key, bs.Size)
}

var BlockCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Interact with raw IPFS blocks.",
		ShortDescription: `
'ipfs block' is a plumbing command used to manipulate raw IPFS blocks.
Reads from stdin or writes to stdout, and <key> is a base58 encoded
multihash.
`,
	},

	Subcommands: map[string]*cmds.Command{
		"stat": blockStatCmd,
		"get":  blockGetCmd,
		"put":  blockPutCmd,
		"rm":   blockRmCmd,
	},
}

var blockStatCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Print information of a raw IPFS block.",
		ShortDescription: `
'ipfs block stat' is a plumbing command for retrieving information
on raw IPFS blocks. It outputs the following to stdout:

	Key  - the base58 encoded multihash
	Size - the size of the block in bytes

`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "The base58 multihash of an existing block to stat.").EnableStdin(),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		b, err := getBlockForKey(req, req.Arguments()[0])
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(&BlockStat{
			Key:  b.Cid().String(),
			Size: len(b.RawData()),
		})
	},
	Type: BlockStat{},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			bs := res.Output().(*BlockStat)
			return strings.NewReader(bs.String()), nil
		},
	},
}

var blockGetCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Get a raw IPFS block.",
		ShortDescription: `
'ipfs block get' is a plumbing command for retrieving raw IPFS blocks.
It outputs to stdout, and <key> is a base58 encoded multihash.
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "The base58 multihash of an existing block to get.").EnableStdin(),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		b, err := getBlockForKey(req, req.Arguments()[0])
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(bytes.NewReader(b.RawData()))
	},
}

var blockPutCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Store input as an IPFS block.",
		ShortDescription: `
'ipfs block put' is a plumbing command for storing raw IPFS blocks.
It reads from stdin, and <key> is a base58 encoded multihash.
`,
	},

	Arguments: []cmds.Argument{
		cmds.FileArg("data", true, false, "The data to be stored as an IPFS block.").EnableStdin(),
	},
	Options: []cmds.Option{
		cmds.StringOption("format", "f", "cid format for blocks to be created with.").Default("v0"),
		cmds.StringOption("mhtype", "multihash hash function").Default("sha2-256"),
		cmds.IntOption("mhlen", "multihash hash length").Default(-1),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		file, err := req.Files().NextFile()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		data, err := ioutil.ReadAll(file)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		err = file.Close()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		var pref cid.Prefix
		pref.Version = 1

		format, _, _ := req.Option("format").String()
		switch format {
		case "cbor":
			pref.Codec = cid.DagCBOR
		case "protobuf":
			pref.Codec = cid.DagProtobuf
		case "raw":
			pref.Codec = cid.Raw
		case "v0":
			pref.Version = 0
			pref.Codec = cid.DagProtobuf
		default:
			res.SetError(fmt.Errorf("unrecognized format: %s", format), cmds.ErrNormal)
			return
		}

		mhtype, _, _ := req.Option("mhtype").String()
		mhtval, ok := mh.Names[mhtype]
		if !ok {
			res.SetError(fmt.Errorf("unrecognized multihash function: %s", mhtype), cmds.ErrNormal)
			return
		}
		pref.MhType = mhtval

		mhlen, _, err := req.Option("mhlen").Int()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}
		pref.MhLength = mhlen

		bcid, err := pref.Sum(data)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		b, err := blocks.NewBlockWithCid(data, bcid)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}
		log.Debugf("BlockPut key: '%q'", b.Cid())

		k, err := n.Blocks.AddBlock(b)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(&BlockStat{
			Key:  k.String(),
			Size: len(data),
		})
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			bs := res.Output().(*BlockStat)
			return strings.NewReader(bs.Key + "\n"), nil
		},
	},
	Type: BlockStat{},
}

func getBlockForKey(req cmds.Request, skey string) (blocks.Block, error) {
	if len(skey) == 0 {
		return nil, fmt.Errorf("zero length cid invalid")
	}

	n, err := req.InvocContext().GetNode()
	if err != nil {
		return nil, err
	}

	c, err := cid.Decode(skey)
	if err != nil {
		return nil, err
	}

	b, err := n.Blocks.GetBlock(req.Context(), c)
	if err != nil {
		return nil, err
	}

	log.Debugf("ipfs block: got block with key: %s", b.Cid())
	return b, nil
}

var blockRmCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Remove IPFS block(s).",
		ShortDescription: `
'ipfs block rm' is a plumbing command for removing raw ipfs blocks.
It takes a list of base58 encoded multihashs to remove.
`,
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("hash", true, true, "Bash58 encoded multihash of block(s) to remove."),
	},
	Options: []cmds.Option{
		cmds.BoolOption("force", "f", "Ignore nonexistent blocks.").Default(false),
		cmds.BoolOption("quiet", "q", "Write minimal output.").Default(false),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}
		hashes := req.Arguments()
		force, _, _ := req.Option("force").Bool()
		quiet, _, _ := req.Option("quiet").Bool()
		cids := make([]*cid.Cid, 0, len(hashes))
		for _, hash := range hashes {
			c, err := cid.Decode(hash)
			if err != nil {
				res.SetError(fmt.Errorf("invalid content id: %s (%s)", hash, err), cmds.ErrNormal)
				return
			}

			cids = append(cids, c)
		}
		ch, err := util.RmBlocks(n.Blockstore, n.Pinning, cids, util.RmBlocksOpts{
			Quiet: quiet,
			Force: force,
		})
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}
		res.SetOutput(ch)
	},
	PostRun: func(req cmds.Request, res cmds.Response) {
		if res.Error() != nil {
			return
		}
		outChan, ok := res.Output().(<-chan interface{})
		if !ok {
			res.SetError(u.ErrCast(), cmds.ErrNormal)
			return
		}
		res.SetOutput(nil)

		err := util.ProcRmOutput(outChan, res.Stdout(), res.Stderr())
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
		}
	},
	Type: util.RemovedBlock{},
}
<commit_msg>block rm: use Marshalers instead of PostRun to process output<commit_after>package commands

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"strings"

	"github.com/ipfs/go-ipfs/blocks"
	util "github.com/ipfs/go-ipfs/blocks/blockstore/util"
	cmds "github.com/ipfs/go-ipfs/commands"

	cid "gx/ipfs/QmV5gPoRsjN1Gid3LMdNZTyfCtP2DsvqEbMAmz82RmmiGk/go-cid"
	u "gx/ipfs/QmZuY8aV7zbNXVy6DyN9SmnuH3o9nG852F4aTiSBpts8d1/go-ipfs-util"
	mh "gx/ipfs/QmbZ6Cee2uHjG7hf19qLHppgKDRtaG4CVtMzdmK9VCVqLu/go-multihash"
)

type BlockStat struct {
	Key  string
	Size int
}

func (bs BlockStat) String() string {
	return fmt.Sprintf("Key: %s\nSize: %d\n", bs.Key, bs.Size)
}

var BlockCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Interact with raw IPFS blocks.",
		ShortDescription: `
'ipfs block' is a plumbing command used to manipulate raw IPFS blocks.
Reads from stdin or writes to stdout, and <key> is a base58 encoded
multihash.
`,
	},

	Subcommands: map[string]*cmds.Command{
		"stat": blockStatCmd,
		"get":  blockGetCmd,
		"put":  blockPutCmd,
		"rm":   blockRmCmd,
	},
}

var blockStatCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Print information of a raw IPFS block.",
		ShortDescription: `
'ipfs block stat' is a plumbing command for retrieving information
on raw IPFS blocks. It outputs the following to stdout:

	Key  - the base58 encoded multihash
	Size - the size of the block in bytes

`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "The base58 multihash of an existing block to stat.").EnableStdin(),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		b, err := getBlockForKey(req, req.Arguments()[0])
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(&BlockStat{
			Key:  b.Cid().String(),
			Size: len(b.RawData()),
		})
	},
	Type: BlockStat{},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			bs := res.Output().(*BlockStat)
			return strings.NewReader(bs.String()), nil
		},
	},
}

var blockGetCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Get a raw IPFS block.",
		ShortDescription: `
'ipfs block get' is a plumbing command for retrieving raw IPFS blocks.
It outputs to stdout, and <key> is a base58 encoded multihash.
`,
	},

	Arguments: []cmds.Argument{
		cmds.StringArg("key", true, false, "The base58 multihash of an existing block to get.").EnableStdin(),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		b, err := getBlockForKey(req, req.Arguments()[0])
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(bytes.NewReader(b.RawData()))
	},
}

var blockPutCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Store input as an IPFS block.",
		ShortDescription: `
'ipfs block put' is a plumbing command for storing raw IPFS blocks.
It reads from stdin, and <key> is a base58 encoded multihash.
`,
	},

	Arguments: []cmds.Argument{
		cmds.FileArg("data", true, false, "The data to be stored as an IPFS block.").EnableStdin(),
	},
	Options: []cmds.Option{
		cmds.StringOption("format", "f", "cid format for blocks to be created with.").Default("v0"),
		cmds.StringOption("mhtype", "multihash hash function").Default("sha2-256"),
		cmds.IntOption("mhlen", "multihash hash length").Default(-1),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		file, err := req.Files().NextFile()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		data, err := ioutil.ReadAll(file)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		err = file.Close()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		var pref cid.Prefix
		pref.Version = 1

		format, _, _ := req.Option("format").String()
		switch format {
		case "cbor":
			pref.Codec = cid.DagCBOR
		case "protobuf":
			pref.Codec = cid.DagProtobuf
		case "raw":
			pref.Codec = cid.Raw
		case "v0":
			pref.Version = 0
			pref.Codec = cid.DagProtobuf
		default:
			res.SetError(fmt.Errorf("unrecognized format: %s", format), cmds.ErrNormal)
			return
		}

		mhtype, _, _ := req.Option("mhtype").String()
		mhtval, ok := mh.Names[mhtype]
		if !ok {
			res.SetError(fmt.Errorf("unrecognized multihash function: %s", mhtype), cmds.ErrNormal)
			return
		}
		pref.MhType = mhtval

		mhlen, _, err := req.Option("mhlen").Int()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}
		pref.MhLength = mhlen

		bcid, err := pref.Sum(data)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		b, err := blocks.NewBlockWithCid(data, bcid)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}
		log.Debugf("BlockPut key: '%q'", b.Cid())

		k, err := n.Blocks.AddBlock(b)
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}

		res.SetOutput(&BlockStat{
			Key:  k.String(),
			Size: len(data),
		})
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			bs := res.Output().(*BlockStat)
			return strings.NewReader(bs.Key + "\n"), nil
		},
	},
	Type: BlockStat{},
}

func getBlockForKey(req cmds.Request, skey string) (blocks.Block, error) {
	if len(skey) == 0 {
		return nil, fmt.Errorf("zero length cid invalid")
	}

	n, err := req.InvocContext().GetNode()
	if err != nil {
		return nil, err
	}

	c, err := cid.Decode(skey)
	if err != nil {
		return nil, err
	}

	b, err := n.Blocks.GetBlock(req.Context(), c)
	if err != nil {
		return nil, err
	}

	log.Debugf("ipfs block: got block with key: %s", b.Cid())
	return b, nil
}

var blockRmCmd = &cmds.Command{
	Helptext: cmds.HelpText{
		Tagline: "Remove IPFS block(s).",
		ShortDescription: `
'ipfs block rm' is a plumbing command for removing raw ipfs blocks.
It takes a list of base58 encoded multihashs to remove.
`,
	},
	Arguments: []cmds.Argument{
		cmds.StringArg("hash", true, true, "Bash58 encoded multihash of block(s) to remove."),
	},
	Options: []cmds.Option{
		cmds.BoolOption("force", "f", "Ignore nonexistent blocks.").Default(false),
		cmds.BoolOption("quiet", "q", "Write minimal output.").Default(false),
	},
	Run: func(req cmds.Request, res cmds.Response) {
		n, err := req.InvocContext().GetNode()
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}
		hashes := req.Arguments()
		force, _, _ := req.Option("force").Bool()
		quiet, _, _ := req.Option("quiet").Bool()
		cids := make([]*cid.Cid, 0, len(hashes))
		for _, hash := range hashes {
			c, err := cid.Decode(hash)
			if err != nil {
				res.SetError(fmt.Errorf("invalid content id: %s (%s)", hash, err), cmds.ErrNormal)
				return
			}

			cids = append(cids, c)
		}
		ch, err := util.RmBlocks(n.Blockstore, n.Pinning, cids, util.RmBlocksOpts{
			Quiet: quiet,
			Force: force,
		})
		if err != nil {
			res.SetError(err, cmds.ErrNormal)
			return
		}
		res.SetOutput(ch)
	},
	Marshalers: cmds.MarshalerMap{
		cmds.Text: func(res cmds.Response) (io.Reader, error) {
			outChan, ok := res.Output().(<-chan interface{})
			if !ok {
				return nil, u.ErrCast()
			}

			err := util.ProcRmOutput(outChan, res.Stdout(), res.Stderr())
			if err != nil {
				return nil, err
			}
			return nil, nil
		},
	},
	Type: util.RemovedBlock{},
}
<|endoftext|>


<commit_before>package pack

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"sync"

	"github.com/restic/restic/internal/debug"
	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/restic"

	"github.com/restic/restic/internal/crypto"
)

// Packer is used to create a new Pack.
type Packer struct {
	blobs []restic.Blob

	bytes uint
	k     *crypto.Key
	wr    io.Writer

	m sync.Mutex
}

// NewPacker returns a new Packer that can be used to pack blobs
// together. If wr is nil, a bytes.Buffer is used.
func NewPacker(k *crypto.Key, wr io.Writer) *Packer {
	if wr == nil {
		wr = bytes.NewBuffer(nil)
	}
	return &Packer{k: k, wr: wr}
}

// Add saves the data read from rd as a new blob to the packer. Returned is the
// number of bytes written to the pack.
func (p *Packer) Add(t restic.BlobType, id restic.ID, data []byte) (int, error) {
	p.m.Lock()
	defer p.m.Unlock()

	c := restic.Blob{Type: t, ID: id}

	n, err := p.wr.Write(data)
	c.Length = uint(n)
	c.Offset = p.bytes
	p.bytes += uint(n)
	p.blobs = append(p.blobs, c)

	return n, errors.Wrap(err, "Write")
}

var entrySize = uint(binary.Size(restic.BlobType(0)) + binary.Size(uint32(0)) + len(restic.ID{}))

// headerEntry is used with encoding/binary to read and write header entries
type headerEntry struct {
	Type   uint8
	Length uint32
	ID     restic.ID
}

// Finalize writes the header for all added blobs and finalizes the pack.
// Returned are the number of bytes written, including the header. If the
// underlying writer implements io.Closer, it is closed.
func (p *Packer) Finalize() (uint, error) {
	p.m.Lock()
	defer p.m.Unlock()

	bytesWritten := p.bytes

	hdrBuf := bytes.NewBuffer(nil)
	bytesHeader, err := p.writeHeader(hdrBuf)
	if err != nil {
		return 0, err
	}

	encryptedHeader := make([]byte, 0, hdrBuf.Len()+p.k.Overhead()+p.k.NonceSize())
	nonce := crypto.NewRandomNonce()
	encryptedHeader = append(encryptedHeader, nonce...)
	encryptedHeader = p.k.Seal(encryptedHeader, nonce, hdrBuf.Bytes(), nil)

	// append the header
	n, err := p.wr.Write(encryptedHeader)
	if err != nil {
		return 0, errors.Wrap(err, "Write")
	}

	hdrBytes := restic.CiphertextLength(int(bytesHeader))
	if n != hdrBytes {
		return 0, errors.New("wrong number of bytes written")
	}

	bytesWritten += uint(hdrBytes)

	// write length
	err = binary.Write(p.wr, binary.LittleEndian, uint32(restic.CiphertextLength(len(p.blobs)*int(entrySize))))
	if err != nil {
		return 0, errors.Wrap(err, "binary.Write")
	}
	bytesWritten += uint(binary.Size(uint32(0)))

	p.bytes = uint(bytesWritten)

	if w, ok := p.wr.(io.Closer); ok {
		return bytesWritten, w.Close()
	}

	return bytesWritten, nil
}

// writeHeader constructs and writes the header to wr.
func (p *Packer) writeHeader(wr io.Writer) (bytesWritten uint, err error) {
	for _, b := range p.blobs {
		entry := headerEntry{
			Length: uint32(b.Length),
			ID:     b.ID,
		}

		switch b.Type {
		case restic.DataBlob:
			entry.Type = 0
		case restic.TreeBlob:
			entry.Type = 1
		default:
			return 0, errors.Errorf("invalid blob type %v", b.Type)
		}

		err := binary.Write(wr, binary.LittleEndian, entry)
		if err != nil {
			return bytesWritten, errors.Wrap(err, "binary.Write")
		}

		bytesWritten += entrySize
	}

	return
}

// Size returns the number of bytes written so far.
func (p *Packer) Size() uint {
	p.m.Lock()
	defer p.m.Unlock()

	return p.bytes
}

// Count returns the number of blobs in this packer.
func (p *Packer) Count() int {
	p.m.Lock()
	defer p.m.Unlock()

	return len(p.blobs)
}

// Blobs returns the slice of blobs that have been written.
func (p *Packer) Blobs() []restic.Blob {
	p.m.Lock()
	defer p.m.Unlock()

	return p.blobs
}

// Writer return the underlying writer.
func (p *Packer) Writer() io.Writer {
	return p.wr
}

func (p *Packer) String() string {
	return fmt.Sprintf("<Packer %d blobs, %d bytes>", len(p.blobs), p.bytes)
}

const maxHeaderSize = 16 * 1024 * 1024

// size of the header-length field at the end of the file
var headerLengthSize = binary.Size(uint32(0))

// we require at least one entry in the header, and one blob for a pack file
var minFileSize = entrySize + crypto.Extension + uint(headerLengthSize)

// number of header enries to download as part of header-length request
var eagerEntries = uint(15)

// readHeader reads the header at the end of rd. size is the length of the
// whole data accessible in rd.
func readHeader(rd io.ReaderAt, size int64) ([]byte, error) {
	debug.Log("size: %v", size)
	if size == 0 {
		err := InvalidFileError{Message: "file is empty"}
		return nil, errors.Wrap(err, "readHeader")
	}

	if size < int64(minFileSize) {
		err := InvalidFileError{Message: "file is too small"}
		return nil, errors.Wrap(err, "readHeader")
	}

	// assuming extra request is significantly slower than extra bytes download,
	// eagerly download eagerEntries header entries as part of header-length request.
	// only make second request if actual number of entries is greater than eagerEntries

	eagerHl := uint32((eagerEntries * entrySize) + crypto.Extension)
	if int64(eagerHl)+int64(headerLengthSize) > size {
		eagerHl = uint32(size) - uint32(headerLengthSize)
	}
	eagerBuf := make([]byte, eagerHl+uint32(headerLengthSize))

	n, err := rd.ReadAt(eagerBuf, size-int64(len(eagerBuf)))
	if err != nil {
		return nil, err
	}
	if n != len(eagerBuf) {
		return nil, errors.New("not enough bytes read")
	}

	hl := binary.LittleEndian.Uint32(eagerBuf[eagerHl:])
	debug.Log("header length: %v", size)

	if hl == 0 {
		err := InvalidFileError{Message: "header length is zero"}
		return nil, errors.Wrap(err, "readHeader")
	}

	if hl < crypto.Extension {
		err := InvalidFileError{Message: "header length is too small"}
		return nil, errors.Wrap(err, "readHeader")
	}

	if (hl-crypto.Extension)%uint32(entrySize) != 0 {
		err := InvalidFileError{Message: "header length is invalid"}
		return nil, errors.Wrap(err, "readHeader")
	}

	if int64(hl) > size-int64(headerLengthSize) {
		err := InvalidFileError{Message: "header is larger than file"}
		return nil, errors.Wrap(err, "readHeader")
	}

	if int64(hl) > maxHeaderSize {
		err := InvalidFileError{Message: "header is larger than maxHeaderSize"}
		return nil, errors.Wrap(err, "readHeader")
	}

	eagerBuf = eagerBuf[:eagerHl]

	var buf []byte
	if hl <= eagerHl {
		// already have all header bytes. yay.
		buf = eagerBuf[eagerHl-hl:]
	} else {
		// need more header bytes
		buf = make([]byte, hl)
		missingHl := hl - eagerHl
		n, err := rd.ReadAt(buf[:missingHl], size-int64(hl)-int64(headerLengthSize))
		if err != nil {
			return nil, errors.Wrap(err, "ReadAt")
		}
		if uint32(n) != missingHl {
			return nil, errors.New("not enough bytes read")
		}
		copy(buf[hl-eagerHl:], eagerBuf)
	}

	return buf, nil
}

// InvalidFileError is return when a file is found that is not a pack file.
type InvalidFileError struct {
	Message string
}

func (e InvalidFileError) Error() string {
	return e.Message
}

// List returns the list of entries found in a pack file.
func List(k *crypto.Key, rd io.ReaderAt, size int64) (entries []restic.Blob, err error) {
	buf, err := readHeader(rd, size)
	if err != nil {
		return nil, err
	}

	if len(buf) < k.NonceSize()+k.Overhead() {
		return nil, errors.New("invalid header, too small")
	}

	nonce, buf := buf[:k.NonceSize()], buf[k.NonceSize():]
	buf, err = k.Open(buf[:0], nonce, buf, nil)
	if err != nil {
		return nil, err
	}

	hdrRd := bytes.NewReader(buf)

	entries = make([]restic.Blob, 0, uint(len(buf))/entrySize)

	pos := uint(0)
	for {
		e := headerEntry{}
		err = binary.Read(hdrRd, binary.LittleEndian, &e)
		if errors.Cause(err) == io.EOF {
			break
		}

		if err != nil {
			return nil, errors.Wrap(err, "binary.Read")
		}

		entry := restic.Blob{
			Length: uint(e.Length),
			ID:     e.ID,
			Offset: pos,
		}

		switch e.Type {
		case 0:
			entry.Type = restic.DataBlob
		case 1:
			entry.Type = restic.TreeBlob
		default:
			return nil, errors.Errorf("invalid type %d", e.Type)
		}

		entries = append(entries, entry)

		pos += uint(e.Length)
	}

	return entries, nil
}
<commit_msg>Refactor the eager-header reads for readability.<commit_after>package pack

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"sync"

	"github.com/restic/restic/internal/debug"
	"github.com/restic/restic/internal/errors"
	"github.com/restic/restic/internal/restic"

	"github.com/restic/restic/internal/crypto"
)

// Packer is used to create a new Pack.
type Packer struct {
	blobs []restic.Blob

	bytes uint
	k     *crypto.Key
	wr    io.Writer

	m sync.Mutex
}

// NewPacker returns a new Packer that can be used to pack blobs
// together. If wr is nil, a bytes.Buffer is used.
func NewPacker(k *crypto.Key, wr io.Writer) *Packer {
	if wr == nil {
		wr = bytes.NewBuffer(nil)
	}
	return &Packer{k: k, wr: wr}
}

// Add saves the data read from rd as a new blob to the packer. Returned is the
// number of bytes written to the pack.
func (p *Packer) Add(t restic.BlobType, id restic.ID, data []byte) (int, error) {
	p.m.Lock()
	defer p.m.Unlock()

	c := restic.Blob{Type: t, ID: id}

	n, err := p.wr.Write(data)
	c.Length = uint(n)
	c.Offset = p.bytes
	p.bytes += uint(n)
	p.blobs = append(p.blobs, c)

	return n, errors.Wrap(err, "Write")
}

var entrySize = uint(binary.Size(restic.BlobType(0)) + binary.Size(uint32(0)) + len(restic.ID{}))

// headerEntry is used with encoding/binary to read and write header entries
type headerEntry struct {
	Type   uint8
	Length uint32
	ID     restic.ID
}

// Finalize writes the header for all added blobs and finalizes the pack.
// Returned are the number of bytes written, including the header. If the
// underlying writer implements io.Closer, it is closed.
func (p *Packer) Finalize() (uint, error) {
	p.m.Lock()
	defer p.m.Unlock()

	bytesWritten := p.bytes

	hdrBuf := bytes.NewBuffer(nil)
	bytesHeader, err := p.writeHeader(hdrBuf)
	if err != nil {
		return 0, err
	}

	encryptedHeader := make([]byte, 0, hdrBuf.Len()+p.k.Overhead()+p.k.NonceSize())
	nonce := crypto.NewRandomNonce()
	encryptedHeader = append(encryptedHeader, nonce...)
	encryptedHeader = p.k.Seal(encryptedHeader, nonce, hdrBuf.Bytes(), nil)

	// append the header
	n, err := p.wr.Write(encryptedHeader)
	if err != nil {
		return 0, errors.Wrap(err, "Write")
	}

	hdrBytes := restic.CiphertextLength(int(bytesHeader))
	if n != hdrBytes {
		return 0, errors.New("wrong number of bytes written")
	}

	bytesWritten += uint(hdrBytes)

	// write length
	err = binary.Write(p.wr, binary.LittleEndian, uint32(restic.CiphertextLength(len(p.blobs)*int(entrySize))))
	if err != nil {
		return 0, errors.Wrap(err, "binary.Write")
	}
	bytesWritten += uint(binary.Size(uint32(0)))

	p.bytes = uint(bytesWritten)

	if w, ok := p.wr.(io.Closer); ok {
		return bytesWritten, w.Close()
	}

	return bytesWritten, nil
}

// writeHeader constructs and writes the header to wr.
func (p *Packer) writeHeader(wr io.Writer) (bytesWritten uint, err error) {
	for _, b := range p.blobs {
		entry := headerEntry{
			Length: uint32(b.Length),
			ID:     b.ID,
		}

		switch b.Type {
		case restic.DataBlob:
			entry.Type = 0
		case restic.TreeBlob:
			entry.Type = 1
		default:
			return 0, errors.Errorf("invalid blob type %v", b.Type)
		}

		err := binary.Write(wr, binary.LittleEndian, entry)
		if err != nil {
			return bytesWritten, errors.Wrap(err, "binary.Write")
		}

		bytesWritten += entrySize
	}

	return
}

// Size returns the number of bytes written so far.
func (p *Packer) Size() uint {
	p.m.Lock()
	defer p.m.Unlock()

	return p.bytes
}

// Count returns the number of blobs in this packer.
func (p *Packer) Count() int {
	p.m.Lock()
	defer p.m.Unlock()

	return len(p.blobs)
}

// Blobs returns the slice of blobs that have been written.
func (p *Packer) Blobs() []restic.Blob {
	p.m.Lock()
	defer p.m.Unlock()

	return p.blobs
}

// Writer return the underlying writer.
func (p *Packer) Writer() io.Writer {
	return p.wr
}

func (p *Packer) String() string {
	return fmt.Sprintf("<Packer %d blobs, %d bytes>", len(p.blobs), p.bytes)
}

var (
	// size of the header-length field at the end of the file
	headerLengthSize = binary.Size(uint32(0))
	// we require at least one entry in the header, and one blob for a pack file
	minFileSize = entrySize + crypto.Extension + uint(headerLengthSize)
)

const (
	maxHeaderSize = 16 * 1024 * 1024
	// number of header enries to download as part of header-length request
	eagerEntries = 15
)

// readRecords reads count records from the underlying ReaderAt, returning the
// raw header, the total number of records in the header, and any error.  If
// the header contains fewer than count entries, the return value is truncated.
func readRecords(rd io.ReaderAt, size int64, count int) ([]byte, int, error) {
	var bufsize int
	bufsize += count * int(entrySize)
	bufsize += crypto.Extension
	bufsize += headerLengthSize

	if bufsize > int(size) {
		bufsize = int(size)
	}

	b := make([]byte, bufsize)
	off := size - int64(bufsize)
	if _, err := rd.ReadAt(b, off); err != nil {
		return nil, 0, err
	}

	header := b[len(b)-headerLengthSize:]
	b = b[:len(b)-headerLengthSize]
	hl := binary.LittleEndian.Uint32(header)
	debug.Log("header length: %v", hl)

	var err error
	switch {
	case hl == 0:
		err = InvalidFileError{Message: "header length is zero"}
	case hl < crypto.Extension:
		err = InvalidFileError{Message: "header length is too small"}
	case (hl-crypto.Extension)%uint32(entrySize) != 0:
		err = InvalidFileError{Message: "header length is invalid"}
	case int64(hl) > size-int64(headerLengthSize):
		err = InvalidFileError{Message: "header is larger than file"}
	case int64(hl) > maxHeaderSize:
		err = InvalidFileError{Message: "header is larger than maxHeaderSize"}
	}
	if err != nil {
		return nil, 0, errors.Wrap(err, "readHeader")
	}

	c := (int(hl) - crypto.Extension) / int(entrySize)
	if c < count {
		recordSize := c * int(entrySize)
		start := len(b) - (recordSize + crypto.Extension)
		b = b[start:]
	}

	return b, c, nil
}

// readHeader reads the header at the end of rd. size is the length of the
// whole data accessible in rd.
func readHeader(rd io.ReaderAt, size int64) ([]byte, error) {
	debug.Log("size: %v", size)
	if size < int64(minFileSize) {
		err := InvalidFileError{Message: "file is too small"}
		return nil, errors.Wrap(err, "readHeader")
	}

	// assuming extra request is significantly slower than extra bytes download,
	// eagerly download eagerEntries header entries as part of header-length request.
	// only make second request if actual number of entries is greater than eagerEntries

	b, c, err := readRecords(rd, size, eagerEntries)
	if err != nil {
		return nil, err
	}
	if c <= eagerEntries {
		// eager read sufficed, return what we got
		return b, nil
	}
	b, _, err = readRecords(rd, size, c)
	if err != nil {
		return nil, err
	}
	return b, nil
}

// InvalidFileError is return when a file is found that is not a pack file.
type InvalidFileError struct {
	Message string
}

func (e InvalidFileError) Error() string {
	return e.Message
}

// List returns the list of entries found in a pack file.
func List(k *crypto.Key, rd io.ReaderAt, size int64) (entries []restic.Blob, err error) {
	buf, err := readHeader(rd, size)
	if err != nil {
		return nil, err
	}

	if len(buf) < k.NonceSize()+k.Overhead() {
		return nil, errors.New("invalid header, too small")
	}

	nonce, buf := buf[:k.NonceSize()], buf[k.NonceSize():]
	buf, err = k.Open(buf[:0], nonce, buf, nil)
	if err != nil {
		return nil, err
	}

	hdrRd := bytes.NewReader(buf)

	entries = make([]restic.Blob, 0, uint(len(buf))/entrySize)

	pos := uint(0)
	for {
		e := headerEntry{}
		err = binary.Read(hdrRd, binary.LittleEndian, &e)
		if errors.Cause(err) == io.EOF {
			break
		}

		if err != nil {
			return nil, errors.Wrap(err, "binary.Read")
		}

		entry := restic.Blob{
			Length: uint(e.Length),
			ID:     e.ID,
			Offset: pos,
		}

		switch e.Type {
		case 0:
			entry.Type = restic.DataBlob
		case 1:
			entry.Type = restic.TreeBlob
		default:
			return nil, errors.Errorf("invalid type %d", e.Type)
		}

		entries = append(entries, entry)

		pos += uint(e.Length)
	}

	return entries, nil
}
<|endoftext|>


<commit_before>package environs_test

import (
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	_ "launchpad.net/juju-core/environs/dummy"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/version"
	"os"
	"path/filepath"
)

type suite struct{}

var _ = Suite(suite{})

var invalidConfigTests = []struct {
	env string
	err string
}{
	{"'", "YAML error:.*"},
	{`
default: unknown
environments:
    only:
        type: unknown
`, `default environment .* does not exist`,
	},
}

func (suite) TestInvalidConfig(c *C) {
	for i, t := range invalidConfigTests {
		c.Logf("running test %v", i)
		_, err := environs.ReadEnvironsBytes([]byte(t.env))
		c.Check(err, ErrorMatches, t.err)
	}
}

var invalidEnvTests = []struct {
	env  string
	name string
	err  string
}{
	{`
environments:
    only:
        foo: bar
`, "", `environment "only" has no type`,
	}, {`
environments:
    only:
        foo: bar
`, "only", `environment "only" has no type`,
	}, {`
environments:
    only:
        foo: bar
        type: crazy
`, "only", `environment "only" has an unknown provider type "crazy"`,
	}, {`
environments:
    only:
        type: dummy
`, "only", `.*state-server: expected bool, got nothing`,
	}, {`
environments:
    only:
        type: dummy
        state-server: false
        unknown-value: causes-an-error
`, "only", `.*unknown-value: expected nothing, got "causes-an-error"`,
	},
}

func (suite) TestInvalidEnv(c *C) {
	defer makeFakeHome(c, "only").restore()
	for i, t := range invalidEnvTests {
		c.Logf("running test %v", i)
		es, err := environs.ReadEnvironsBytes([]byte(t.env))
		c.Check(err, IsNil)
		e, err := es.Open(t.name)
		c.Check(err, ErrorMatches, t.err)
		c.Check(e, IsNil)
	}
}

var configTests = []struct {
	env   string
	check func(c *C, es *environs.Environs)
}{
	{`
environments:
    only:
        type: dummy
        state-server: false
`, func(c *C, es *environs.Environs) {
		e, err := es.Open("")
		c.Assert(err, IsNil)
		c.Assert(e.Name(), Equals, "only")
	}}, {`
default:
    invalid
environments:
    valid:
        type: dummy
        state-server: false
    invalid:
        type: crazy
`, func(c *C, es *environs.Environs) {
		e, err := es.Open("")
		c.Assert(err, ErrorMatches, `environment "invalid" has an unknown provider type "crazy"`)
		c.Assert(e, IsNil)
		e, err = es.Open("valid")
		c.Assert(err, IsNil)
		c.Assert(e.Name(), Equals, "valid")
	}}, {`
environments:
    one:
        type: dummy
        state-server: false
    two:
        type: dummy
        state-server: false
`, func(c *C, es *environs.Environs) {
		e, err := es.Open("")
		c.Assert(err, ErrorMatches, `no default environment found`)
		c.Assert(e, IsNil)
	}},
}

func (suite) TestConfig(c *C) {
	defer makeFakeHome(c, "only", "valid", "one", "two").restore()
	for i, t := range configTests {
		c.Logf("running test %v", i)
		es, err := environs.ReadEnvironsBytes([]byte(t.env))
		c.Assert(err, IsNil)
		t.check(c, es)
	}
}

func (suite) TestDefaultConfigFile(c *C) {
	defer makeFakeHome(c, "only").restore()

	env := `
environments:
    only:
        type: dummy
        state-server: false
        authorized-keys: i-am-a-key
`
	err := ioutil.WriteFile(homePath(".juju", "environments.yaml"), []byte(env), 0666)
	c.Assert(err, IsNil)

	es, err := environs.ReadEnvirons("")
	c.Assert(err, IsNil)
	e, err := es.Open("")
	c.Assert(err, IsNil)
	c.Assert(e.Name(), Equals, "only")
}

func (suite) TestNamedConfigFile(c *C) {
	defer makeFakeHome(c, "only").restore()

	env := `
environments:
    only:
        type: dummy
        state-server: false
        authorized-keys: i-am-a-key
`
	path := filepath.Join(c.MkDir(), "a-file")
	err := ioutil.WriteFile(path, []byte(env), 0666)
	c.Assert(err, IsNil)

	es, err := environs.ReadEnvirons(path)
	c.Assert(err, IsNil)
	e, err := es.Open("")
	c.Assert(err, IsNil)
	c.Assert(e.Name(), Equals, "only")
}

func (suite) TestConfigRoundTrip(c *C) {
	cfg, err := config.New(map[string]interface{}{
		"name":            "bladaam",
		"type":            "dummy",
		"state-server":    false,
		"authorized-keys": "i-am-a-key",
		"ca-cert":         testing.CACertPEM,
		"ca-private-key":  "",
	})
	c.Assert(err, IsNil)
	provider, err := environs.Provider(cfg.Type())
	c.Assert(err, IsNil)
	cfg, err = provider.Validate(cfg, nil)
	c.Assert(err, IsNil)
	env, err := environs.New(cfg)
	c.Assert(err, IsNil)
	c.Assert(cfg.AllAttrs(), DeepEquals, env.Config().AllAttrs())
}

func (suite) TestBootstrapConfig(c *C) {
	defer makeFakeHome(c, "bladaam").restore()
	cfg, err := config.New(map[string]interface{}{
		"name":            "bladaam",
		"type":            "dummy",
		"state-server":    false,
		"admin-secret":    "highly",
		"secret":          "um",
		"authorized-keys": "i-am-a-key",
		"ca-cert":         testing.CACertPEM,
		"ca-private-key":  testing.CAKeyPEM,
	})
	c.Assert(err, IsNil)
	provider, err := environs.Provider(cfg.Type())
	c.Assert(err, IsNil)

	tools := &state.Tools{
		URL:    "http://x",
		Binary: version.MustParseBinary("1.2.3-foo-bar"),
	}
	cfg1, err := environs.BootstrapConfig(provider, cfg, tools)
	c.Assert(err, IsNil)

	expect := cfg.AllAttrs()
	delete(expect, "secret")
	expect["admin-secret"] = ""
	expect["ca-private-key"] = nil
	expect["agent-version"] = "1.2.3"
	c.Assert(cfg1.AllAttrs(), DeepEquals, expect)
}

type fakeHome string

func makeFakeHome(c *C, certNames ...string) fakeHome {
	oldHome := os.Getenv("HOME")
	os.Setenv("HOME", c.MkDir())

	err := os.Mkdir(homePath(".juju"), 0777)
	c.Assert(err, IsNil)
	for _, name := range certNames {
		err := ioutil.WriteFile(homePath(".juju", name+"-cert.pem"), []byte(testing.CACertPEM), 0666)
		c.Assert(err, IsNil)
		err = ioutil.WriteFile(homePath(".juju", name+"-private-key.pem"), []byte(testing.CAKeyPEM), 0666)
		c.Assert(err, IsNil)
	}

	err = os.Mkdir(homePath(".ssh"), 0777)
	c.Assert(err, IsNil)
	err = ioutil.WriteFile(homePath(".ssh", "id_rsa.pub"), []byte("auth key\n"), 0666)
	c.Assert(err, IsNil)

	return fakeHome(oldHome)
}

func homePath(names ...string) string {
	all := append([]string{os.Getenv("HOME")}, names...)
	return filepath.Join(all...)
}

func (h fakeHome) restore() {
	os.Setenv("HOME", string(h))
}
<commit_msg>environs: fix config test<commit_after>package environs_test

import (
	"io/ioutil"
	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/config"
	_ "launchpad.net/juju-core/environs/dummy"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/version"
	"os"
	"path/filepath"
)

type suite struct{}

var _ = Suite(suite{})

var invalidConfigTests = []struct {
	env string
	err string
}{
	{"'", "YAML error:.*"},
	{`
default: unknown
environments:
    only:
        type: unknown
`, `default environment .* does not exist`,
	},
}

func (suite) TestInvalidConfig(c *C) {
	for i, t := range invalidConfigTests {
		c.Logf("running test %v", i)
		_, err := environs.ReadEnvironsBytes([]byte(t.env))
		c.Check(err, ErrorMatches, t.err)
	}
}

var invalidEnvTests = []struct {
	env  string
	name string
	err  string
}{
	{`
environments:
    only:
        foo: bar
`, "", `environment "only" has no type`,
	}, {`
environments:
    only:
        foo: bar
`, "only", `environment "only" has no type`,
	}, {`
environments:
    only:
        foo: bar
        type: crazy
`, "only", `environment "only" has an unknown provider type "crazy"`,
	}, {`
environments:
    only:
        type: dummy
`, "only", `.*state-server: expected bool, got nothing`,
	}, {`
environments:
    only:
        type: dummy
        state-server: false
        unknown-value: causes-an-error
`, "only", `.*unknown-value: expected nothing, got "causes-an-error"`,
	},
}

func (suite) TestInvalidEnv(c *C) {
	defer makeFakeHome(c, "only").restore()
	for i, t := range invalidEnvTests {
		c.Logf("running test %v", i)
		es, err := environs.ReadEnvironsBytes([]byte(t.env))
		c.Check(err, IsNil)
		e, err := es.Open(t.name)
		c.Check(err, ErrorMatches, t.err)
		c.Check(e, IsNil)
	}
}

var configTests = []struct {
	env   string
	check func(c *C, es *environs.Environs)
}{
	{`
environments:
    only:
        type: dummy
        state-server: false
`, func(c *C, es *environs.Environs) {
		e, err := es.Open("")
		c.Assert(err, IsNil)
		c.Assert(e.Name(), Equals, "only")
	}}, {`
default:
    invalid
environments:
    valid:
        type: dummy
        state-server: false
    invalid:
        type: crazy
`, func(c *C, es *environs.Environs) {
		e, err := es.Open("")
		c.Assert(err, ErrorMatches, `environment "invalid" has an unknown provider type "crazy"`)
		c.Assert(e, IsNil)
		e, err = es.Open("valid")
		c.Assert(err, IsNil)
		c.Assert(e.Name(), Equals, "valid")
	}}, {`
environments:
    one:
        type: dummy
        state-server: false
    two:
        type: dummy
        state-server: false
`, func(c *C, es *environs.Environs) {
		e, err := es.Open("")
		c.Assert(err, ErrorMatches, `no default environment found`)
		c.Assert(e, IsNil)
	}},
}

func (suite) TestConfig(c *C) {
	defer makeFakeHome(c, "only", "valid", "one", "two").restore()
	for i, t := range configTests {
		c.Logf("running test %v", i)
		es, err := environs.ReadEnvironsBytes([]byte(t.env))
		c.Assert(err, IsNil)
		t.check(c, es)
	}
}

func (suite) TestDefaultConfigFile(c *C) {
	defer makeFakeHome(c, "only").restore()

	env := `
environments:
    only:
        type: dummy
        state-server: false
        authorized-keys: i-am-a-key
`
	err := ioutil.WriteFile(homePath(".juju", "environments.yaml"), []byte(env), 0666)
	c.Assert(err, IsNil)

	es, err := environs.ReadEnvirons("")
	c.Assert(err, IsNil)
	e, err := es.Open("")
	c.Assert(err, IsNil)
	c.Assert(e.Name(), Equals, "only")
}

func (suite) TestNamedConfigFile(c *C) {
	defer makeFakeHome(c, "only").restore()

	env := `
environments:
    only:
        type: dummy
        state-server: false
        authorized-keys: i-am-a-key
`
	path := filepath.Join(c.MkDir(), "a-file")
	err := ioutil.WriteFile(path, []byte(env), 0666)
	c.Assert(err, IsNil)

	es, err := environs.ReadEnvirons(path)
	c.Assert(err, IsNil)
	e, err := es.Open("")
	c.Assert(err, IsNil)
	c.Assert(e.Name(), Equals, "only")
}

func (suite) TestConfigRoundTrip(c *C) {
	cfg, err := config.New(map[string]interface{}{
		"name":            "bladaam",
		"type":            "dummy",
		"state-server":    false,
		"authorized-keys": "i-am-a-key",
		"ca-cert":         testing.CACertPEM,
		"ca-private-key":  "",
	})
	c.Assert(err, IsNil)
	provider, err := environs.Provider(cfg.Type())
	c.Assert(err, IsNil)
	cfg, err = provider.Validate(cfg, nil)
	c.Assert(err, IsNil)
	env, err := environs.New(cfg)
	c.Assert(err, IsNil)
	c.Assert(cfg.AllAttrs(), DeepEquals, env.Config().AllAttrs())
}

func (suite) TestBootstrapConfig(c *C) {
	defer makeFakeHome(c, "bladaam").restore()
	cfg, err := config.New(map[string]interface{}{
		"name":            "bladaam",
		"type":            "dummy",
		"state-server":    false,
		"admin-secret":    "highly",
		"secret":          "um",
		"authorized-keys": "i-am-a-key",
		"ca-cert":         testing.CACertPEM,
		"ca-private-key":  testing.CAKeyPEM,
	})
	c.Assert(err, IsNil)
	provider, err := environs.Provider(cfg.Type())
	c.Assert(err, IsNil)

	tools := &state.Tools{
		URL:    "http://x",
		Binary: version.MustParseBinary("1.2.3-foo-bar"),
	}
	cfg1, err := environs.BootstrapConfig(provider, cfg, tools)
	c.Assert(err, IsNil)

	expect := cfg.AllAttrs()
	delete(expect, "secret")
	expect["admin-secret"] = ""
	expect["ca-private-key"] = ""
	expect["agent-version"] = "1.2.3"
	c.Assert(cfg1.AllAttrs(), DeepEquals, expect)
}

type fakeHome string

func makeFakeHome(c *C, certNames ...string) fakeHome {
	oldHome := os.Getenv("HOME")
	os.Setenv("HOME", c.MkDir())

	err := os.Mkdir(homePath(".juju"), 0777)
	c.Assert(err, IsNil)
	for _, name := range certNames {
		err := ioutil.WriteFile(homePath(".juju", name+"-cert.pem"), []byte(testing.CACertPEM), 0666)
		c.Assert(err, IsNil)
		err = ioutil.WriteFile(homePath(".juju", name+"-private-key.pem"), []byte(testing.CAKeyPEM), 0666)
		c.Assert(err, IsNil)
	}

	err = os.Mkdir(homePath(".ssh"), 0777)
	c.Assert(err, IsNil)
	err = ioutil.WriteFile(homePath(".ssh", "id_rsa.pub"), []byte("auth key\n"), 0666)
	c.Assert(err, IsNil)

	return fakeHome(oldHome)
}

func homePath(names ...string) string {
	all := append([]string{os.Getenv("HOME")}, names...)
	return filepath.Join(all...)
}

func (h fakeHome) restore() {
	os.Setenv("HOME", string(h))
}
<|endoftext|>


<commit_before>package pool

import (
	"errors"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"gopkg.in/bsm/ratelimit.v1"

	"gopkg.in/redis.v4/internal"
)

var (
	ErrClosed      = errors.New("redis: client is closed")
	ErrPoolTimeout = errors.New("redis: connection pool timeout")
	errConnStale   = errors.New("connection is stale")
)

var timers = sync.Pool{
	New: func() interface{} {
		return time.NewTimer(0)
	},
}

// PoolStats contains pool state information and accumulated stats.
type PoolStats struct {
	Requests uint32 // number of times a connection was requested by the pool
	Hits     uint32 // number of times free connection was found in the pool
	Timeouts uint32 // number of times a wait timeout occurred

	TotalConns uint32 // the number of total connections in the pool
	FreeConns  uint32 // the number of free connections in the pool
}

type Pooler interface {
	Get() (*Conn, error)
	Put(*Conn) error
	Remove(*Conn, error) error
	Len() int
	FreeLen() int
	Stats() *PoolStats
	Close() error
	Closed() bool
}

type dialer func() (net.Conn, error)

type ConnPool struct {
	_dial       dialer
	DialLimiter *ratelimit.RateLimiter
	OnClose     func(*Conn) error

	poolTimeout time.Duration
	idleTimeout time.Duration

	queue chan struct{}

	connsMu sync.Mutex
	conns   []*Conn

	freeConnsMu sync.Mutex
	freeConns   []*Conn

	stats PoolStats

	_closed int32 // atomic
	lastErr atomic.Value
}

var _ Pooler = (*ConnPool)(nil)

func NewConnPool(dial dialer, poolSize int, poolTimeout, idleTimeout, idleCheckFrequency time.Duration) *ConnPool {
	p := &ConnPool{
		_dial:       dial,
		DialLimiter: ratelimit.New(3*poolSize, time.Second),

		poolTimeout: poolTimeout,
		idleTimeout: idleTimeout,

		queue:     make(chan struct{}, poolSize),
		conns:     make([]*Conn, 0, poolSize),
		freeConns: make([]*Conn, 0, poolSize),
	}
	for i := 0; i < poolSize; i++ {
		p.queue <- struct{}{}
	}
	if idleTimeout > 0 && idleCheckFrequency > 0 {
		go p.reaper(idleCheckFrequency)
	}
	return p
}

func (p *ConnPool) dial() (net.Conn, error) {
	if p.DialLimiter != nil && p.DialLimiter.Limit() {
		err := fmt.Errorf(
			"redis: you open connections too fast (last_error=%q)",
			p.loadLastErr(),
		)
		return nil, err
	}

	cn, err := p._dial()
	if err != nil {
		p.storeLastErr(err.Error())
		return nil, err
	}
	return cn, nil
}

func (p *ConnPool) NewConn() (*Conn, error) {
	netConn, err := p.dial()
	if err != nil {
		return nil, err
	}
	return NewConn(netConn), nil
}

func (p *ConnPool) PopFree() *Conn {
	timer := timers.Get().(*time.Timer)
	if !timer.Reset(p.poolTimeout) {
		<-timer.C
	}

	select {
	case <-p.queue:
		timers.Put(timer)
	case <-timer.C:
		timers.Put(timer)
		atomic.AddUint32(&p.stats.Timeouts, 1)
		return nil
	}

	p.freeConnsMu.Lock()
	cn := p.popFree()
	p.freeConnsMu.Unlock()

	if cn == nil {
		p.queue <- struct{}{}
	}
	return cn
}

func (p *ConnPool) popFree() *Conn {
	if len(p.freeConns) == 0 {
		return nil
	}

	idx := len(p.freeConns) - 1
	cn := p.freeConns[idx]
	p.freeConns = p.freeConns[:idx]
	return cn
}

// Get returns existed connection from the pool or creates a new one.
func (p *ConnPool) Get() (*Conn, error) {
	if p.Closed() {
		return nil, ErrClosed
	}

	atomic.AddUint32(&p.stats.Requests, 1)

	timer := timers.Get().(*time.Timer)
	if !timer.Reset(p.poolTimeout) {
		<-timer.C
	}

	select {
	case <-p.queue:
		timers.Put(timer)
	case <-timer.C:
		timers.Put(timer)
		atomic.AddUint32(&p.stats.Timeouts, 1)
		return nil, ErrPoolTimeout
	}

	p.freeConnsMu.Lock()
	cn := p.popFree()
	p.freeConnsMu.Unlock()

	if cn != nil {
		atomic.AddUint32(&p.stats.Hits, 1)
		if !cn.IsStale(p.idleTimeout) {
			return cn, nil
		}
		_ = p.closeConn(cn, errConnStale)
	}

	newcn, err := p.NewConn()
	if err != nil {
		p.queue <- struct{}{}
		return nil, err
	}

	p.connsMu.Lock()
	if cn != nil {
		p.removeConn(cn)
	}
	p.conns = append(p.conns, newcn)
	p.connsMu.Unlock()

	return newcn, nil
}

func (p *ConnPool) Put(cn *Conn) error {
	if data := cn.Rd.PeekBuffered(); data != nil {
		err := fmt.Errorf("connection has unread data: %q", data)
		internal.Logf(err.Error())
		return p.Remove(cn, err)
	}
	p.freeConnsMu.Lock()
	p.freeConns = append(p.freeConns, cn)
	p.freeConnsMu.Unlock()
	p.queue <- struct{}{}
	return nil
}

func (p *ConnPool) Remove(cn *Conn, reason error) error {
	p.remove(cn, reason)
	p.queue <- struct{}{}
	return nil
}

func (p *ConnPool) remove(cn *Conn, reason error) {
	_ = p.closeConn(cn, reason)

	p.connsMu.Lock()
	p.removeConn(cn)
	p.connsMu.Unlock()
}

func (p *ConnPool) removeConn(cn *Conn) {
	for i, c := range p.conns {
		if c == cn {
			p.conns = append(p.conns[:i], p.conns[i+1:]...)
			break
		}
	}
}

// Len returns total number of connections.
func (p *ConnPool) Len() int {
	p.connsMu.Lock()
	l := len(p.conns)
	p.connsMu.Unlock()
	return l
}

// FreeLen returns number of free connections.
func (p *ConnPool) FreeLen() int {
	p.freeConnsMu.Lock()
	l := len(p.freeConns)
	p.freeConnsMu.Unlock()
	return l
}

func (p *ConnPool) Stats() *PoolStats {
	stats := PoolStats{}
	stats.Requests = atomic.LoadUint32(&p.stats.Requests)
	stats.Hits = atomic.LoadUint32(&p.stats.Hits)
	stats.Timeouts = atomic.LoadUint32(&p.stats.Timeouts)
	stats.TotalConns = uint32(p.Len())
	stats.FreeConns = uint32(p.FreeLen())
	return &stats
}

func (p *ConnPool) Closed() bool {
	return atomic.LoadInt32(&p._closed) == 1
}

func (p *ConnPool) Close() (retErr error) {
	if !atomic.CompareAndSwapInt32(&p._closed, 0, 1) {
		return ErrClosed
	}

	p.connsMu.Lock()
	// Close all connections.
	for _, cn := range p.conns {
		if cn == nil {
			continue
		}
		if err := p.closeConn(cn, ErrClosed); err != nil && retErr == nil {
			retErr = err
		}
	}
	p.conns = nil
	p.connsMu.Unlock()

	p.freeConnsMu.Lock()
	p.freeConns = nil
	p.freeConnsMu.Unlock()

	return retErr
}

func (p *ConnPool) closeConn(cn *Conn, reason error) error {
	p.storeLastErr(reason.Error())
	if p.OnClose != nil {
		_ = p.OnClose(cn)
	}
	return cn.Close()
}

func (p *ConnPool) reapStaleConn() bool {
	if len(p.freeConns) == 0 {
		return false
	}

	cn := p.freeConns[0]
	if !cn.IsStale(p.idleTimeout) {
		return false
	}

	p.remove(cn, errConnStale)
	p.freeConns = append(p.freeConns[:0], p.freeConns[1:]...)

	return true
}

func (p *ConnPool) ReapStaleConns() (int, error) {
	var n int
	for {
		<-p.queue
		p.freeConnsMu.Lock()

		reaped := p.reapStaleConn()

		p.freeConnsMu.Unlock()
		p.queue <- struct{}{}

		if reaped {
			n++
		} else {
			break
		}
	}
	return n, nil
}

func (p *ConnPool) reaper(frequency time.Duration) {
	ticker := time.NewTicker(frequency)
	defer ticker.Stop()

	for _ = range ticker.C {
		if p.Closed() {
			break
		}
		n, err := p.ReapStaleConns()
		if err != nil {
			internal.Logf("ReapStaleConns failed: %s", err)
			continue
		}
		s := p.Stats()
		internal.Logf(
			"reaper: removed %d stale conns (TotalConns=%d FreeConns=%d Requests=%d Hits=%d Timeouts=%d)",
			n, s.TotalConns, s.FreeConns, s.Requests, s.Hits, s.Timeouts,
		)
	}
}

func (p *ConnPool) storeLastErr(err string) {
	p.lastErr.Store(err)
}

func (p *ConnPool) loadLastErr() string {
	if v := p.lastErr.Load(); v != nil {
		return v.(string)
	}
	return ""
}

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

var idleCheckFrequency atomic.Value

func SetIdleCheckFrequency(d time.Duration) {
	idleCheckFrequency.Store(d)
}

func getIdleCheckFrequency() time.Duration {
	v := idleCheckFrequency.Load()
	if v == nil {
		return time.Minute
	}
	return v.(time.Duration)
}
<commit_msg>internal/pool: more idiomatic work with channels.<commit_after>package pool

import (
	"errors"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"gopkg.in/bsm/ratelimit.v1"

	"gopkg.in/redis.v4/internal"
)

var (
	ErrClosed      = errors.New("redis: client is closed")
	ErrPoolTimeout = errors.New("redis: connection pool timeout")
	errConnStale   = errors.New("connection is stale")
)

var timers = sync.Pool{
	New: func() interface{} {
		return time.NewTimer(0)
	},
}

// PoolStats contains pool state information and accumulated stats.
type PoolStats struct {
	Requests uint32 // number of times a connection was requested by the pool
	Hits     uint32 // number of times free connection was found in the pool
	Timeouts uint32 // number of times a wait timeout occurred

	TotalConns uint32 // the number of total connections in the pool
	FreeConns  uint32 // the number of free connections in the pool
}

type Pooler interface {
	Get() (*Conn, error)
	Put(*Conn) error
	Remove(*Conn, error) error
	Len() int
	FreeLen() int
	Stats() *PoolStats
	Close() error
	Closed() bool
}

type dialer func() (net.Conn, error)

type ConnPool struct {
	_dial       dialer
	DialLimiter *ratelimit.RateLimiter
	OnClose     func(*Conn) error

	poolTimeout time.Duration
	idleTimeout time.Duration

	queue chan struct{}

	connsMu sync.Mutex
	conns   []*Conn

	freeConnsMu sync.Mutex
	freeConns   []*Conn

	stats PoolStats

	_closed int32 // atomic
	lastErr atomic.Value
}

var _ Pooler = (*ConnPool)(nil)

func NewConnPool(dial dialer, poolSize int, poolTimeout, idleTimeout, idleCheckFrequency time.Duration) *ConnPool {
	p := &ConnPool{
		_dial:       dial,
		DialLimiter: ratelimit.New(3*poolSize, time.Second),

		poolTimeout: poolTimeout,
		idleTimeout: idleTimeout,

		queue:     make(chan struct{}, poolSize),
		conns:     make([]*Conn, 0, poolSize),
		freeConns: make([]*Conn, 0, poolSize),
	}
	if idleTimeout > 0 && idleCheckFrequency > 0 {
		go p.reaper(idleCheckFrequency)
	}
	return p
}

func (p *ConnPool) dial() (net.Conn, error) {
	if p.DialLimiter != nil && p.DialLimiter.Limit() {
		err := fmt.Errorf(
			"redis: you open connections too fast (last_error=%q)",
			p.loadLastErr(),
		)
		return nil, err
	}

	cn, err := p._dial()
	if err != nil {
		p.storeLastErr(err.Error())
		return nil, err
	}
	return cn, nil
}

func (p *ConnPool) NewConn() (*Conn, error) {
	netConn, err := p.dial()
	if err != nil {
		return nil, err
	}
	return NewConn(netConn), nil
}

func (p *ConnPool) PopFree() *Conn {
	timer := timers.Get().(*time.Timer)
	if !timer.Reset(p.poolTimeout) {
		<-timer.C
	}

	select {
	case p.queue <- struct{}{}:
		timers.Put(timer)
	case <-timer.C:
		timers.Put(timer)
		atomic.AddUint32(&p.stats.Timeouts, 1)
		return nil
	}

	p.freeConnsMu.Lock()
	cn := p.popFree()
	p.freeConnsMu.Unlock()

	if cn == nil {
		<-p.queue
	}
	return cn
}

func (p *ConnPool) popFree() *Conn {
	if len(p.freeConns) == 0 {
		return nil
	}

	idx := len(p.freeConns) - 1
	cn := p.freeConns[idx]
	p.freeConns = p.freeConns[:idx]
	return cn
}

// Get returns existed connection from the pool or creates a new one.
func (p *ConnPool) Get() (*Conn, error) {
	if p.Closed() {
		return nil, ErrClosed
	}

	atomic.AddUint32(&p.stats.Requests, 1)

	timer := timers.Get().(*time.Timer)
	if !timer.Reset(p.poolTimeout) {
		<-timer.C
	}

	select {
	case p.queue <- struct{}{}:
		timers.Put(timer)
	case <-timer.C:
		timers.Put(timer)
		atomic.AddUint32(&p.stats.Timeouts, 1)
		return nil, ErrPoolTimeout
	}

	p.freeConnsMu.Lock()
	cn := p.popFree()
	p.freeConnsMu.Unlock()

	if cn != nil {
		atomic.AddUint32(&p.stats.Hits, 1)
		if !cn.IsStale(p.idleTimeout) {
			return cn, nil
		}
		_ = p.closeConn(cn, errConnStale)
	}

	newcn, err := p.NewConn()
	if err != nil {
		<-p.queue
		return nil, err
	}

	p.connsMu.Lock()
	if cn != nil {
		p.removeConn(cn)
	}
	p.conns = append(p.conns, newcn)
	p.connsMu.Unlock()

	return newcn, nil
}

func (p *ConnPool) Put(cn *Conn) error {
	if data := cn.Rd.PeekBuffered(); data != nil {
		err := fmt.Errorf("connection has unread data: %q", data)
		internal.Logf(err.Error())
		return p.Remove(cn, err)
	}
	p.freeConnsMu.Lock()
	p.freeConns = append(p.freeConns, cn)
	p.freeConnsMu.Unlock()
	<-p.queue
	return nil
}

func (p *ConnPool) Remove(cn *Conn, reason error) error {
	p.remove(cn, reason)
	<-p.queue
	return nil
}

func (p *ConnPool) remove(cn *Conn, reason error) {
	_ = p.closeConn(cn, reason)

	p.connsMu.Lock()
	p.removeConn(cn)
	p.connsMu.Unlock()
}

func (p *ConnPool) removeConn(cn *Conn) {
	for i, c := range p.conns {
		if c == cn {
			p.conns = append(p.conns[:i], p.conns[i+1:]...)
			break
		}
	}
}

// Len returns total number of connections.
func (p *ConnPool) Len() int {
	p.connsMu.Lock()
	l := len(p.conns)
	p.connsMu.Unlock()
	return l
}

// FreeLen returns number of free connections.
func (p *ConnPool) FreeLen() int {
	p.freeConnsMu.Lock()
	l := len(p.freeConns)
	p.freeConnsMu.Unlock()
	return l
}

func (p *ConnPool) Stats() *PoolStats {
	stats := PoolStats{}
	stats.Requests = atomic.LoadUint32(&p.stats.Requests)
	stats.Hits = atomic.LoadUint32(&p.stats.Hits)
	stats.Timeouts = atomic.LoadUint32(&p.stats.Timeouts)
	stats.TotalConns = uint32(p.Len())
	stats.FreeConns = uint32(p.FreeLen())
	return &stats
}

func (p *ConnPool) Closed() bool {
	return atomic.LoadInt32(&p._closed) == 1
}

func (p *ConnPool) Close() (retErr error) {
	if !atomic.CompareAndSwapInt32(&p._closed, 0, 1) {
		return ErrClosed
	}

	p.connsMu.Lock()
	// Close all connections.
	for _, cn := range p.conns {
		if cn == nil {
			continue
		}
		if err := p.closeConn(cn, ErrClosed); err != nil && retErr == nil {
			retErr = err
		}
	}
	p.conns = nil
	p.connsMu.Unlock()

	p.freeConnsMu.Lock()
	p.freeConns = nil
	p.freeConnsMu.Unlock()

	return retErr
}

func (p *ConnPool) closeConn(cn *Conn, reason error) error {
	p.storeLastErr(reason.Error())
	if p.OnClose != nil {
		_ = p.OnClose(cn)
	}
	return cn.Close()
}

func (p *ConnPool) reapStaleConn() bool {
	if len(p.freeConns) == 0 {
		return false
	}

	cn := p.freeConns[0]
	if !cn.IsStale(p.idleTimeout) {
		return false
	}

	p.remove(cn, errConnStale)
	p.freeConns = append(p.freeConns[:0], p.freeConns[1:]...)

	return true
}

func (p *ConnPool) ReapStaleConns() (int, error) {
	var n int
	for {
		p.queue <- struct{}{}
		p.freeConnsMu.Lock()

		reaped := p.reapStaleConn()

		p.freeConnsMu.Unlock()
		<-p.queue

		if reaped {
			n++
		} else {
			break
		}
	}
	return n, nil
}

func (p *ConnPool) reaper(frequency time.Duration) {
	ticker := time.NewTicker(frequency)
	defer ticker.Stop()

	for _ = range ticker.C {
		if p.Closed() {
			break
		}
		n, err := p.ReapStaleConns()
		if err != nil {
			internal.Logf("ReapStaleConns failed: %s", err)
			continue
		}
		s := p.Stats()
		internal.Logf(
			"reaper: removed %d stale conns (TotalConns=%d FreeConns=%d Requests=%d Hits=%d Timeouts=%d)",
			n, s.TotalConns, s.FreeConns, s.Requests, s.Hits, s.Timeouts,
		)
	}
}

func (p *ConnPool) storeLastErr(err string) {
	p.lastErr.Store(err)
}

func (p *ConnPool) loadLastErr() string {
	if v := p.lastErr.Load(); v != nil {
		return v.(string)
	}
	return ""
}

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

var idleCheckFrequency atomic.Value

func SetIdleCheckFrequency(d time.Duration) {
	idleCheckFrequency.Store(d)
}

func getIdleCheckFrequency() time.Duration {
	v := idleCheckFrequency.Load()
	if v == nil {
		return time.Minute
	}
	return v.(time.Duration)
}
<|endoftext|>