<commit_before>// Package cache_server contains core functionality for our cache servers; storing & retrieving files etc.
package server

import (
	"bytes"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"core"
)

// A cachedFile stores metadata about a file stored in our cache.
type cachedFile struct {
	// Arbitrates single access to this file
	sync.RWMutex
	// Time the file was last read
	lastReadTime time.Time
	// Number of times the file has been read
	readCount int
	// Size of the file
	size int64
}

// Handles synchronisation of storing/retrieving artifacts.
var cachedFiles = map[string]*cachedFile{}
var mutex sync.RWMutex
var totalSize int64

var cachePath = ".plz-cache"

// Init initialises the cache and fires off a background cleaner goroutine which runs every
// cleanFrequency seconds. The high and low water marks control a (soft) max size and a (harder)
// minimum size.
// This causes an initial scan of the directory and is needed to retrieve any preexisting artifacts.
// If it's not called the cache will still function but won't know about anything existing beforehand.
func Init(path string, cleanFrequency int, lowWaterMark, highWaterMark int64) {
	log.Notice("Initialising cache with settings:\n  Path: %s\n  Clean frequency: %d\n  Low water mark: %d\n  High water mark: %d",
		path, cleanFrequency, lowWaterMark, highWaterMark)
	cachePath = path
	scan(path)
	go clean(path, cleanFrequency, lowWaterMark, highWaterMark)
}

// scan scans the directory tree for files.
func scan(path string) {
	mutex.Lock()
	defer mutex.Unlock()
	cachedFiles = map[string]*cachedFile{}
	totalSize = 0

	if !core.PathExists(path) {
		if err := os.MkdirAll(path, core.DirPermissions); err != nil {
			log.Fatalf("Failed to create cache directory %s: %s", path, err)
		}
		return
	}

	log.Info("Scanning cache directory %s...", path)
	filepath.Walk(path, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			log.Fatalf("%s", err)
		} else if !info.IsDir() { // We don't have directory entries.
			name = name[len(path)+1 : len(name)]
			log.Debug("Found file %s", name)
			size := info.Size()
			cachedFiles[name] = &cachedFile{
				lastReadTime: time.Now(),
				readCount:    0,
				size:         size,
			}
			totalSize += size
		}
		return nil
	})
	cachePath = path
	log.Info("Scan complete")
}

// lockFile locks a file for reading or writing.
// It returns a locked mutex corresponding to that file or nil if there is none.
// The caller should .Unlock() the mutex once they're done with it.
func lockFile(path string, write bool, size int64) *cachedFile {
	mutex.Lock()
	defer mutex.Unlock()
	file, present := cachedFiles[path]
	if !present {
		// If we're writing we insert a new one, if we're reading we don't.
		if !write {
			return nil
		}
		file = &cachedFile{
			lastReadTime: time.Now(),
			readCount:    0,
			size:         size,
		}
		cachedFiles[path] = file
		totalSize += size
	}
	if write {
		file.Lock()
	} else {
		file.RLock()
		file.readCount++
	}
	return file
}

// removeFile deletes a file from the cache map. It does not remove the on-disk file.
// Caller should already hold the mutex before calling this
func removeFile(path string, file *cachedFile) {
	delete(cachedFiles, path)
	totalSize -= file.size
	log.Debug("Removing file %s, saves %d, new size will be %d", path, file.size, totalSize)
}

// RetrieveArtifact takes in the artifact path as a parameter and checks in the base server
// file directory to see if the file exists in the given path. If found, the function will
// return whatever's been stored there, which might be a directory and therefore contain
// multiple files to be returned.
func RetrieveArtifact(artPath string) (map[string][]byte, error) {
	fullPath := path.Join(cachePath, artPath)
	lock := lockFile(artPath, false, 0)
	if lock == nil {
		// Can happen if artPath is a directory; we only store artifacts as files.
		// (This is a debatable choice; it's a bit crap either way).
		if info, err := os.Stat(fullPath); err == nil && info.IsDir() {
			return retrieveDir(artPath)
		}
		return nil, os.ErrNotExist
	}
	defer lock.RUnlock()

	ret := map[string][]byte{}
	if err := filepath.Walk(fullPath, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			body, err := ioutil.ReadFile(name)
			if err != nil {
				return err
			}
			ret[name[len(cachePath)+1:len(name)]] = body
		}
		return nil
	}); err != nil {
		return nil, err
	}
	return ret, nil
}

// retrieveDir retrieves a directory of artifacts. We don't track the directory itself
// but allow its traversal to retrieve them.
func retrieveDir(artPath string) (map[string][]byte, error) {
	log.Debug("Searching dir %s for artifacts", artPath)
	ret := map[string][]byte{}
	fullPath := path.Join(cachePath, artPath)
	err := filepath.Walk(fullPath, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			// Must strip cache path off the front of this.
			m, err := RetrieveArtifact(name[len(cachePath)+1:])
			if err != nil {
				return err
			}
			for k, v := range m {
				ret[k] = v
			}
		}
		return nil
	})
	return ret, err
}

// StoreArtifact takes in the artifact content and path as parameters and creates a file with
// the given content in the given path.
// The function will return the first error found in the process, or nil if the process is successful.
func StoreArtifact(artPath string, key []byte) error {
	log.Info("Storing artifact %s", artPath)
	lock := lockFile(artPath, true, int64(len(key)))
	defer lock.Unlock()

	artPath = path.Join(cachePath, artPath)
	dirPath := path.Dir(artPath)
	if err := os.MkdirAll(dirPath, core.DirPermissions); err != nil {
		log.Warning("Couldn't create path %s in http cache: %s", dirPath, err)
		return err
	}
	log.Debug("Writing artifact to %s", artPath)
	if err := core.WriteFile(bytes.NewReader(key), artPath, 0); err != nil {
		log.Errorf("Could not create %s artifact: %s", artPath, err)
		return err
	}
	return nil
}

// DeleteArtifact takes in the artifact path as a parameter and removes the artifact from disk.
// The function will return the first error found in the process, or nil if the process is successful.
func DeleteArtifact(artPath string) error {
	// We need to search the entire map for prefixes. Pessimism follows...
	mutex.Lock()
	for p, file := range cachedFiles {
		if strings.HasPrefix(p, artPath) {
			file.Lock()
			removeFile(path.Join(cachePath, p), file)
			file.Unlock()
		}
	}
	mutex.Unlock()
	return os.RemoveAll(path.Join(cachePath, artPath))
}

// DeleteAllArtifacts will remove all files in the cache directory.
// The function will return the first error found in the process, or nil if the process is successful.
func DeleteAllArtifacts() error {
	// Empty entire cache now.
	mutex.Lock()
	cachedFiles = map[string]*cachedFile{}
	mutex.Unlock()
	defer scan(cachePath) // Rescan whatever's left afterwards.

	files, err := ioutil.ReadDir(cachePath)
	if err != nil && os.IsNotExist(err) {
		log.Warning("%s directory does not exist, nothing to delete.", cachePath)
		return nil
	} else if err != nil {
		log.Errorf("Error reading cache directory: %s", err)
		return err
	}
	for _, file := range files {
		p := path.Join(cachePath, file.Name())
		if err := os.RemoveAll(p); err != nil {
			log.Errorf("Failed to remove directory %s: %s", p, err)
			return err
		}
	}
	return nil
}

// clean implements a periodic clean of the cache to remove old artifacts.
func clean(cachePath string, cleanFrequency int, lowWaterMark, highWaterMark int64) {
	for _ = range time.NewTicker(time.Duration(cleanFrequency) * time.Second).C {
		log.Debug("Total size: %d High water mark: %d", totalSize, highWaterMark)
		if totalSize > highWaterMark {
			log.Info("Cleaning cache...")
			files := filesToClean(lowWaterMark)
			log.Info("Identified %d files to clean...", len(files))
			for _, file := range files {
				if err := os.RemoveAll(path.Join(cachePath, file.path)); err != nil {
					log.Errorf("Failed to remove artifact: %s", err)
				} else {
					removeFile(file.path, file.file)
				}
			}
		}
	}
}

// cachedFilePath embeds a cachedFile but with the path too.
type cachedFilePath struct {
	file *cachedFile
	path string
}

type cachedFilePaths []cachedFilePath

func (c cachedFilePaths) Len() int      { return len(c) }
func (c cachedFilePaths) Swap(i, j int) { c[i], c[j] = c[j], c[i] }
func (c cachedFilePaths) Less(i, j int) bool {
	return c[i].file.lastReadTime.Before(c[j].file.lastReadTime)
}

// filesToClean returns a list of files that should be cleaned, ie. the least interesting
// artifacts in the cache according to some heuristic. Removing all of them will be
// sufficient to reduce the cache size below lowWaterMark.
func filesToClean(lowWaterMark int64) cachedFilePaths {
	mutex.Lock()
	ret := make(cachedFilePaths, 0, len(cachedFiles))
	for path, file := range cachedFiles {
		ret = append(ret, cachedFilePath{file: file, path: path})
	}
	mutex.Unlock()
	sort.Sort(&ret)

	sizeToDelete := totalSize - lowWaterMark
	var sizeDeleted int64
	for i, file := range ret {
		if sizeDeleted >= sizeToDelete {
			return ret[0:i]
		}
		sizeDeleted += file.file.size
	}
	return ret
}
<commit_msg>Have cache be more careful about deleting artifacts if they don't write correctly.<commit_after>// Package cache_server contains core functionality for our cache servers; storing & retrieving files etc.
package server

import (
	"bytes"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"core"
)

// A cachedFile stores metadata about a file stored in our cache.
type cachedFile struct {
	// Arbitrates single access to this file
	sync.RWMutex
	// Time the file was last read
	lastReadTime time.Time
	// Number of times the file has been read
	readCount int
	// Size of the file
	size int64
}

// Handles synchronisation of storing/retrieving artifacts.
var cachedFiles = map[string]*cachedFile{}
var mutex sync.RWMutex
var totalSize int64

var cachePath = ".plz-cache"

// Init initialises the cache and fires off a background cleaner goroutine which runs every
// cleanFrequency seconds. The high and low water marks control a (soft) max size and a (harder)
// minimum size.
// This causes an initial scan of the directory and is needed to retrieve any preexisting artifacts.
// If it's not called the cache will still function but won't know about anything existing beforehand.
func Init(path string, cleanFrequency int, lowWaterMark, highWaterMark int64) {
	log.Notice("Initialising cache with settings:\n  Path: %s\n  Clean frequency: %d\n  Low water mark: %d\n  High water mark: %d",
		path, cleanFrequency, lowWaterMark, highWaterMark)
	cachePath = path
	scan(path)
	go clean(path, cleanFrequency, lowWaterMark, highWaterMark)
}

// scan scans the directory tree for files.
func scan(path string) {
	mutex.Lock()
	defer mutex.Unlock()
	cachedFiles = map[string]*cachedFile{}
	totalSize = 0

	if !core.PathExists(path) {
		if err := os.MkdirAll(path, core.DirPermissions); err != nil {
			log.Fatalf("Failed to create cache directory %s: %s", path, err)
		}
		return
	}

	log.Info("Scanning cache directory %s...", path)
	filepath.Walk(path, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			log.Fatalf("%s", err)
		} else if !info.IsDir() { // We don't have directory entries.
			name = name[len(path)+1 : len(name)]
			log.Debug("Found file %s", name)
			size := info.Size()
			cachedFiles[name] = &cachedFile{
				lastReadTime: time.Now(),
				readCount:    0,
				size:         size,
			}
			totalSize += size
		}
		return nil
	})
	cachePath = path
	log.Info("Scan complete")
}

// lockFile locks a file for reading or writing.
// It returns a locked mutex corresponding to that file or nil if there is none.
// The caller should .Unlock() the mutex once they're done with it.
func lockFile(path string, write bool, size int64) *cachedFile {
	mutex.Lock()
	defer mutex.Unlock()
	file, present := cachedFiles[path]
	if !present {
		// If we're writing we insert a new one, if we're reading we don't.
		if !write {
			return nil
		}
		file = &cachedFile{
			lastReadTime: time.Now(),
			readCount:    0,
			size:         size,
		}
		cachedFiles[path] = file
		totalSize += size
	}
	if write {
		file.Lock()
	} else {
		file.RLock()
		file.readCount++
	}
	return file
}

// removeFile deletes a file from the cache map. It does not remove the on-disk file.
// Caller should already hold the mutex before calling this
func removeFile(path string, file *cachedFile) {
	delete(cachedFiles, path)
	totalSize -= file.size
	log.Debug("Removing file %s, saves %d, new size will be %d", path, file.size, totalSize)
}

// removeAndDeleteFile deletes a file from the cache map and on-disk.
// Caller should already hold the mutex before calling this
func removeAndDeleteFile(p string, file *cachedFile) {
	removeFile(p, file)
	p = path.Join(cachePath, p)
	if err := os.RemoveAll(p); err != nil {
		log.Error("Failed to delete file: %s", p)
	}
}

// RetrieveArtifact takes in the artifact path as a parameter and checks in the base server
// file directory to see if the file exists in the given path. If found, the function will
// return whatever's been stored there, which might be a directory and therefore contain
// multiple files to be returned.
func RetrieveArtifact(artPath string) (map[string][]byte, error) {
	fullPath := path.Join(cachePath, artPath)
	lock := lockFile(artPath, false, 0)
	if lock == nil {
		// Can happen if artPath is a directory; we only store artifacts as files.
		// (This is a debatable choice; it's a bit crap either way).
		if info, err := os.Stat(fullPath); err == nil && info.IsDir() {
			return retrieveDir(artPath)
		}
		return nil, os.ErrNotExist
	}
	defer lock.RUnlock()

	ret := map[string][]byte{}
	if err := filepath.Walk(fullPath, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			body, err := ioutil.ReadFile(name)
			if err != nil {
				return err
			}
			ret[name[len(cachePath)+1:len(name)]] = body
		}
		return nil
	}); err != nil {
		return nil, err
	}
	return ret, nil
}

// retrieveDir retrieves a directory of artifacts. We don't track the directory itself
// but allow its traversal to retrieve them.
func retrieveDir(artPath string) (map[string][]byte, error) {
	log.Debug("Searching dir %s for artifacts", artPath)
	ret := map[string][]byte{}
	fullPath := path.Join(cachePath, artPath)
	err := filepath.Walk(fullPath, func(name string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		} else if !info.IsDir() {
			// Must strip cache path off the front of this.
			m, err := RetrieveArtifact(name[len(cachePath)+1:])
			if err != nil {
				return err
			}
			for k, v := range m {
				ret[k] = v
			}
		}
		return nil
	})
	return ret, err
}

// StoreArtifact takes in the artifact content and path as parameters and creates a file with
// the given content in the given path.
// The function will return the first error found in the process, or nil if the process is successful.
func StoreArtifact(artPath string, key []byte) error {
	log.Info("Storing artifact %s", artPath)
	lock := lockFile(artPath, true, int64(len(key)))
	defer lock.Unlock()

	fullPath := path.Join(cachePath, artPath)
	dirPath := path.Dir(fullPath)
	if err := os.MkdirAll(dirPath, core.DirPermissions); err != nil {
		log.Warning("Couldn't create path %s in http cache: %s", dirPath, err)
		removeAndDeleteFile(artPath, lock)
		return err
	}
	log.Debug("Writing artifact to %s", fullPath)
	if err := core.WriteFile(bytes.NewReader(key), fullPath, 0); err != nil {
		log.Errorf("Could not create %s artifact: %s", fullPath, err)
		removeAndDeleteFile(artPath, lock)
		return err
	}
	return nil
}

// DeleteArtifact takes in the artifact path as a parameter and removes the artifact from disk.
// The function will return the first error found in the process, or nil if the process is successful.
func DeleteArtifact(artPath string) error {
	// We need to search the entire map for prefixes. Pessimism follows...
	mutex.Lock()
	for p, file := range cachedFiles {
		if strings.HasPrefix(p, artPath) {
			file.Lock()
			removeFile(p, file)
			file.Unlock()
		}
	}
	mutex.Unlock()
	return os.RemoveAll(path.Join(cachePath, artPath))
}

// DeleteAllArtifacts will remove all files in the cache directory.
// The function will return the first error found in the process, or nil if the process is successful.
func DeleteAllArtifacts() error {
	// Empty entire cache now.
	mutex.Lock()
	cachedFiles = map[string]*cachedFile{}
	mutex.Unlock()
	defer scan(cachePath) // Rescan whatever's left afterwards.

	files, err := ioutil.ReadDir(cachePath)
	if err != nil && os.IsNotExist(err) {
		log.Warning("%s directory does not exist, nothing to delete.", cachePath)
		return nil
	} else if err != nil {
		log.Errorf("Error reading cache directory: %s", err)
		return err
	}
	for _, file := range files {
		p := path.Join(cachePath, file.Name())
		if err := os.RemoveAll(p); err != nil {
			log.Errorf("Failed to remove directory %s: %s", p, err)
			return err
		}
	}
	return nil
}

// clean implements a periodic clean of the cache to remove old artifacts.
func clean(cachePath string, cleanFrequency int, lowWaterMark, highWaterMark int64) {
	for _ = range time.NewTicker(time.Duration(cleanFrequency) * time.Second).C {
		log.Debug("Total size: %d High water mark: %d", totalSize, highWaterMark)
		if totalSize > highWaterMark {
			log.Info("Cleaning cache...")
			files := filesToClean(lowWaterMark)
			log.Info("Identified %d files to clean...", len(files))
			for _, file := range files {
				if err := os.RemoveAll(path.Join(cachePath, file.path)); err != nil {
					log.Errorf("Failed to remove artifact: %s", err)
				} else {
					removeFile(file.path, file.file)
				}
			}
		}
	}
}

// cachedFilePath embeds a cachedFile but with the path too.
type cachedFilePath struct {
	file *cachedFile
	path string
}

type cachedFilePaths []cachedFilePath

func (c cachedFilePaths) Len() int      { return len(c) }
func (c cachedFilePaths) Swap(i, j int) { c[i], c[j] = c[j], c[i] }
func (c cachedFilePaths) Less(i, j int) bool {
	return c[i].file.lastReadTime.Before(c[j].file.lastReadTime)
}

// filesToClean returns a list of files that should be cleaned, ie. the least interesting
// artifacts in the cache according to some heuristic. Removing all of them will be
// sufficient to reduce the cache size below lowWaterMark.
func filesToClean(lowWaterMark int64) cachedFilePaths {
	mutex.Lock()
	ret := make(cachedFilePaths, 0, len(cachedFiles))
	for path, file := range cachedFiles {
		ret = append(ret, cachedFilePath{file: file, path: path})
	}
	mutex.Unlock()
	sort.Sort(&ret)

	sizeToDelete := totalSize - lowWaterMark
	var sizeDeleted int64
	for i, file := range ret {
		if sizeDeleted >= sizeToDelete {
			return ret[0:i]
		}
		sizeDeleted += file.file.size
	}
	return ret
}
<|endoftext|>


<commit_before>package kingpin

//go:generate go run ./cmd/genvalues/main.go

import (
	"fmt"
	"net"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"strings"
	"time"

	"github.com/alecthomas/units"
)

// NOTE: Most of the base type values were lifted from:
// http://golang.org/src/pkg/flag/flag.go?s=20146:20222

// Value is the interface to the dynamic value stored in a flag.
// (The default value is represented as a string.)
//
// If a Value has an IsBoolFlag() bool method returning true, the command-line
// parser makes --name equivalent to -name=true rather than using the next
// command-line argument, and adds a --no-name counterpart for negating the
// flag.
type Value interface {
	String() string
	Set(string) error
	Get() interface{}
}

// Optional interface to indicate boolean flags that don't accept a value, and
// implicitly have a --no-<x> negation counterpart.
type boolFlag interface {
	Value
	IsBoolFlag() bool
}

// Optional interface for arguments that cumulatively consume all remaining
// input.
type remainderArg interface {
	Value
	IsCumulative() bool
}

type accumulator struct {
	element func(value interface{}) Value
	typ     reflect.Type
	slice   reflect.Value
}

// Use reflection to accumulate values into a slice.
//
// target := []string{}
// newAccumulator(&target, func (value interface{}) Value {
//   return newStringValue(value.(*string))
// })
func newAccumulator(slice interface{}, element func(value interface{}) Value) *accumulator {
	typ := reflect.TypeOf(slice)
	if typ.Kind() != reflect.Ptr || typ.Elem().Kind() != reflect.Slice {
		panic("expected a pointer to a slice")
	}
	return &accumulator{
		element: element,
		typ:     typ.Elem().Elem(),
		slice:   reflect.ValueOf(slice),
	}
}

func (a *accumulator) String() string {
	out := []string{}
	s := a.slice.Elem()
	for i := 0; i < s.Len(); i++ {
		out = append(out, a.element(s.Index(i).Addr().Interface()).String())
	}
	return strings.Join(out, ",")
}

func (a *accumulator) Set(value string) error {
	e := reflect.New(a.typ)
	if err := a.element(e.Interface()).Set(value); err != nil {
		return err
	}
	slice := reflect.Append(a.slice.Elem(), e.Elem())
	a.slice.Elem().Set(slice)
	return nil
}

func (a *accumulator) Get() interface{} {
	return a.slice.Interface()
}

func (a *accumulator) IsCumulative() bool {
	return true
}

func (b *boolValue) IsBoolFlag() bool { return true }

// -- time.Duration Value
type durationValue time.Duration

func newDurationValue(p *time.Duration) *durationValue {
	return (*durationValue)(p)
}

func (d *durationValue) Set(s string) error {
	v, err := time.ParseDuration(s)
	*d = durationValue(v)
	return err
}

func (d *durationValue) Get() interface{} { return time.Duration(*d) }

func (d *durationValue) String() string { return (*time.Duration)(d).String() }

// -- map[string]string Value
type stringMapValue map[string]string

func newStringMapValue(p *map[string]string) *stringMapValue {
	return (*stringMapValue)(p)
}

var stringMapRegex = regexp.MustCompile("[:=]")

func (s *stringMapValue) Set(value string) error {
	parts := stringMapRegex.Split(value, 2)
	if len(parts) != 2 {
		return fmt.Errorf("expected KEY=VALUE got '%s'", value)
	}
	(*s)[parts[0]] = parts[1]
	return nil
}

func (s *stringMapValue) Get() interface{} {
	return (map[string]string)(*s)
}

func (s *stringMapValue) String() string {
	return fmt.Sprintf("%s", map[string]string(*s))
}

func (s *stringMapValue) IsCumulative() bool {
	return true
}

// -- net.IP Value
type ipValue net.IP

func newIPValue(p *net.IP) *ipValue {
	return (*ipValue)(p)
}

func (i *ipValue) Set(value string) error {
	if ip := net.ParseIP(value); ip == nil {
		return fmt.Errorf("'%s' is not an IP address", value)
	} else {
		*i = *(*ipValue)(&ip)
		return nil
	}
}

func (i *ipValue) Get() interface{} {
	return (net.IP)(*i)
}

func (i *ipValue) String() string {
	return (*net.IP)(i).String()
}

// -- *net.TCPAddr Value
type tcpAddrValue struct {
	addr **net.TCPAddr
}

func newTCPAddrValue(p **net.TCPAddr) *tcpAddrValue {
	return &tcpAddrValue{p}
}

func (i *tcpAddrValue) Set(value string) error {
	if addr, err := net.ResolveTCPAddr("tcp", value); err != nil {
		return fmt.Errorf("'%s' is not a valid TCP address: %s", value, err)
	} else {
		*i.addr = addr
		return nil
	}
}

func (t *tcpAddrValue) Get() interface{} {
	return (*net.TCPAddr)(*t.addr)
}

func (i *tcpAddrValue) String() string {
	return (*i.addr).String()
}

// -- existingFile Value

type fileStatValue struct {
	path      *string
	predicate func(os.FileInfo) error
}

func newFileStatValue(p *string, predicate func(os.FileInfo) error) *fileStatValue {
	return &fileStatValue{
		path:      p,
		predicate: predicate,
	}
}

func (e *fileStatValue) Set(value string) error {
	if s, err := os.Stat(value); os.IsNotExist(err) {
		return fmt.Errorf("path '%s' does not exist", value)
	} else if err != nil {
		return err
	} else if err := e.predicate(s); err != nil {
		return err
	}
	*e.path = value
	return nil
}

func (f *fileStatValue) Get() interface{} {
	return (string)(*f.path)
}

func (e *fileStatValue) String() string {
	return *e.path
}

// -- os.File value

type fileValue struct {
	f    **os.File
	flag int
	perm os.FileMode
}

func newFileValue(p **os.File, flag int, perm os.FileMode) *fileValue {
	return &fileValue{p, flag, perm}
}

func (f *fileValue) Set(value string) error {
	if fd, err := os.OpenFile(value, f.flag, f.perm); err != nil {
		return err
	} else {
		*f.f = fd
		return nil
	}
}

func (f *fileValue) Get() interface{} {
	return (*os.File)(*f.f)
}

func (f *fileValue) String() string {
	if *f.f == nil {
		return "<nil>"
	}
	return (*f.f).Name()
}

// -- url.URL Value
type urlValue struct {
	u **url.URL
}

func newURLValue(p **url.URL) *urlValue {
	return &urlValue{p}
}

func (u *urlValue) Set(value string) error {
	if url, err := url.Parse(value); err != nil {
		return fmt.Errorf("invalid URL: %s", err)
	} else {
		*u.u = url
		return nil
	}
}

func (u *urlValue) Get() interface{} {
	return (*url.URL)(*u.u)
}

func (u *urlValue) String() string {
	if *u.u == nil {
		return "<nil>"
	}
	return (*u.u).String()
}

// -- []*url.URL Value
type urlListValue []*url.URL

func newURLListValue(p *[]*url.URL) *urlListValue {
	return (*urlListValue)(p)
}

func (u *urlListValue) Set(value string) error {
	if url, err := url.Parse(value); err != nil {
		return fmt.Errorf("invalid URL: %s", err)
	} else {
		*u = append(*u, url)
		return nil
	}
}

func (u *urlListValue) Get() interface{} {
	return ([]*url.URL)(*u)
}

func (u *urlListValue) String() string {
	out := []string{}
	for _, url := range *u {
		out = append(out, url.String())
	}
	return strings.Join(out, ",")
}

// A flag whose value must be in a set of options.
type enumValue struct {
	value   *string
	options []string
}

func newEnumFlag(target *string, options ...string) *enumValue {
	return &enumValue{
		value:   target,
		options: options,
	}
}

func (a *enumValue) String() string {
	return *a.value
}

func (a *enumValue) Set(value string) error {
	for _, v := range a.options {
		if v == value {
			*a.value = value
			return nil
		}
	}
	return fmt.Errorf("enum value must be one of %s, got '%s'", strings.Join(a.options, ","), value)
}

func (e *enumValue) Get() interface{} {
	return (string)(*e.value)
}

// -- []string Enum Value
type enumsValue struct {
	value   *[]string
	options []string
}

func newEnumsFlag(target *[]string, options ...string) *enumsValue {
	return &enumsValue{
		value:   target,
		options: options,
	}
}

func (s *enumsValue) Set(value string) error {
	for _, v := range s.options {
		if v == value {
			*s.value = append(*s.value, value)
			return nil
		}
	}
	return fmt.Errorf("enum value must be one of %s, got '%s'", strings.Join(s.options, ","), value)
}

func (e *enumsValue) Get() interface{} {
	return ([]string)(*e.value)
}

func (s *enumsValue) String() string {
	return strings.Join(*s.value, ",")
}

func (s *enumsValue) IsCumulative() bool {
	return true
}

// -- units.Base2Bytes Value
type bytesValue units.Base2Bytes

func newBytesValue(p *units.Base2Bytes) *bytesValue {
	return (*bytesValue)(p)
}

func (d *bytesValue) Set(s string) error {
	v, err := units.ParseBase2Bytes(s)
	*d = bytesValue(v)
	return err
}

func (d *bytesValue) Get() interface{} { return units.Base2Bytes(*d) }

func (d *bytesValue) String() string { return (*units.Base2Bytes)(d).String() }

func newExistingFileValue(target *string) *fileStatValue {
	return newFileStatValue(target, func(s os.FileInfo) error {
		if s.IsDir() {
			return fmt.Errorf("'%s' is a directory", s.Name())
		}
		return nil
	})
}

func newExistingDirValue(target *string) *fileStatValue {
	return newFileStatValue(target, func(s os.FileInfo) error {
		if !s.IsDir() {
			return fmt.Errorf("'%s' is a file", s.Name())
		}
		return nil
	})
}

func newExistingFileOrDirValue(target *string) *fileStatValue {
	return newFileStatValue(target, func(s os.FileInfo) error { return nil })
}

type counterValue int

func newCounterValue(n *int) *counterValue {
	return (*counterValue)(n)
}

func (c *counterValue) Set(s string) error {
	*c++
	return nil
}

func (c *counterValue) Get() interface{} { return (int)(*c) }
func (c *counterValue) IsBoolFlag() bool { return true }
func (c *counterValue) String() string   { return fmt.Sprintf("%d", *c) }
<commit_msg>Pull Get out of Value again to avoid breaking external values.<commit_after>package kingpin

//go:generate go run ./cmd/genvalues/main.go

import (
	"fmt"
	"net"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"strings"
	"time"

	"github.com/alecthomas/units"
)

// NOTE: Most of the base type values were lifted from:
// http://golang.org/src/pkg/flag/flag.go?s=20146:20222

// Value is the interface to the dynamic value stored in a flag.
// (The default value is represented as a string.)
//
// If a Value has an IsBoolFlag() bool method returning true, the command-line
// parser makes --name equivalent to -name=true rather than using the next
// command-line argument, and adds a --no-name counterpart for negating the
// flag.
type Value interface {
	String() string
	Set(string) error
}

// Getter is an interface that allows the contents of a Value to be retrieved.
// It wraps the Value interface, rather than being part of it, because it
// appeared after Go 1 and its compatibility rules. All Value types provided
// by this package satisfy the Getter interface.
type Getter interface {
	Value
	Get() interface{}
}

// Optional interface to indicate boolean flags that don't accept a value, and
// implicitly have a --no-<x> negation counterpart.
type boolFlag interface {
	Value
	IsBoolFlag() bool
}

// Optional interface for arguments that cumulatively consume all remaining
// input.
type remainderArg interface {
	Value
	IsCumulative() bool
}

type accumulator struct {
	element func(value interface{}) Value
	typ     reflect.Type
	slice   reflect.Value
}

// Use reflection to accumulate values into a slice.
//
// target := []string{}
// newAccumulator(&target, func (value interface{}) Value {
//   return newStringValue(value.(*string))
// })
func newAccumulator(slice interface{}, element func(value interface{}) Value) *accumulator {
	typ := reflect.TypeOf(slice)
	if typ.Kind() != reflect.Ptr || typ.Elem().Kind() != reflect.Slice {
		panic("expected a pointer to a slice")
	}
	return &accumulator{
		element: element,
		typ:     typ.Elem().Elem(),
		slice:   reflect.ValueOf(slice),
	}
}

func (a *accumulator) String() string {
	out := []string{}
	s := a.slice.Elem()
	for i := 0; i < s.Len(); i++ {
		out = append(out, a.element(s.Index(i).Addr().Interface()).String())
	}
	return strings.Join(out, ",")
}

func (a *accumulator) Set(value string) error {
	e := reflect.New(a.typ)
	if err := a.element(e.Interface()).Set(value); err != nil {
		return err
	}
	slice := reflect.Append(a.slice.Elem(), e.Elem())
	a.slice.Elem().Set(slice)
	return nil
}

func (a *accumulator) Get() interface{} {
	return a.slice.Interface()
}

func (a *accumulator) IsCumulative() bool {
	return true
}

func (b *boolValue) IsBoolFlag() bool { return true }

// -- time.Duration Value
type durationValue time.Duration

func newDurationValue(p *time.Duration) *durationValue {
	return (*durationValue)(p)
}

func (d *durationValue) Set(s string) error {
	v, err := time.ParseDuration(s)
	*d = durationValue(v)
	return err
}

func (d *durationValue) Get() interface{} { return time.Duration(*d) }

func (d *durationValue) String() string { return (*time.Duration)(d).String() }

// -- map[string]string Value
type stringMapValue map[string]string

func newStringMapValue(p *map[string]string) *stringMapValue {
	return (*stringMapValue)(p)
}

var stringMapRegex = regexp.MustCompile("[:=]")

func (s *stringMapValue) Set(value string) error {
	parts := stringMapRegex.Split(value, 2)
	if len(parts) != 2 {
		return fmt.Errorf("expected KEY=VALUE got '%s'", value)
	}
	(*s)[parts[0]] = parts[1]
	return nil
}

func (s *stringMapValue) Get() interface{} {
	return (map[string]string)(*s)
}

func (s *stringMapValue) String() string {
	return fmt.Sprintf("%s", map[string]string(*s))
}

func (s *stringMapValue) IsCumulative() bool {
	return true
}

// -- net.IP Value
type ipValue net.IP

func newIPValue(p *net.IP) *ipValue {
	return (*ipValue)(p)
}

func (i *ipValue) Set(value string) error {
	if ip := net.ParseIP(value); ip == nil {
		return fmt.Errorf("'%s' is not an IP address", value)
	} else {
		*i = *(*ipValue)(&ip)
		return nil
	}
}

func (i *ipValue) Get() interface{} {
	return (net.IP)(*i)
}

func (i *ipValue) String() string {
	return (*net.IP)(i).String()
}

// -- *net.TCPAddr Value
type tcpAddrValue struct {
	addr **net.TCPAddr
}

func newTCPAddrValue(p **net.TCPAddr) *tcpAddrValue {
	return &tcpAddrValue{p}
}

func (i *tcpAddrValue) Set(value string) error {
	if addr, err := net.ResolveTCPAddr("tcp", value); err != nil {
		return fmt.Errorf("'%s' is not a valid TCP address: %s", value, err)
	} else {
		*i.addr = addr
		return nil
	}
}

func (t *tcpAddrValue) Get() interface{} {
	return (*net.TCPAddr)(*t.addr)
}

func (i *tcpAddrValue) String() string {
	return (*i.addr).String()
}

// -- existingFile Value

type fileStatValue struct {
	path      *string
	predicate func(os.FileInfo) error
}

func newFileStatValue(p *string, predicate func(os.FileInfo) error) *fileStatValue {
	return &fileStatValue{
		path:      p,
		predicate: predicate,
	}
}

func (e *fileStatValue) Set(value string) error {
	if s, err := os.Stat(value); os.IsNotExist(err) {
		return fmt.Errorf("path '%s' does not exist", value)
	} else if err != nil {
		return err
	} else if err := e.predicate(s); err != nil {
		return err
	}
	*e.path = value
	return nil
}

func (f *fileStatValue) Get() interface{} {
	return (string)(*f.path)
}

func (e *fileStatValue) String() string {
	return *e.path
}

// -- os.File value

type fileValue struct {
	f    **os.File
	flag int
	perm os.FileMode
}

func newFileValue(p **os.File, flag int, perm os.FileMode) *fileValue {
	return &fileValue{p, flag, perm}
}

func (f *fileValue) Set(value string) error {
	if fd, err := os.OpenFile(value, f.flag, f.perm); err != nil {
		return err
	} else {
		*f.f = fd
		return nil
	}
}

func (f *fileValue) Get() interface{} {
	return (*os.File)(*f.f)
}

func (f *fileValue) String() string {
	if *f.f == nil {
		return "<nil>"
	}
	return (*f.f).Name()
}

// -- url.URL Value
type urlValue struct {
	u **url.URL
}

func newURLValue(p **url.URL) *urlValue {
	return &urlValue{p}
}

func (u *urlValue) Set(value string) error {
	if url, err := url.Parse(value); err != nil {
		return fmt.Errorf("invalid URL: %s", err)
	} else {
		*u.u = url
		return nil
	}
}

func (u *urlValue) Get() interface{} {
	return (*url.URL)(*u.u)
}

func (u *urlValue) String() string {
	if *u.u == nil {
		return "<nil>"
	}
	return (*u.u).String()
}

// -- []*url.URL Value
type urlListValue []*url.URL

func newURLListValue(p *[]*url.URL) *urlListValue {
	return (*urlListValue)(p)
}

func (u *urlListValue) Set(value string) error {
	if url, err := url.Parse(value); err != nil {
		return fmt.Errorf("invalid URL: %s", err)
	} else {
		*u = append(*u, url)
		return nil
	}
}

func (u *urlListValue) Get() interface{} {
	return ([]*url.URL)(*u)
}

func (u *urlListValue) String() string {
	out := []string{}
	for _, url := range *u {
		out = append(out, url.String())
	}
	return strings.Join(out, ",")
}

// A flag whose value must be in a set of options.
type enumValue struct {
	value   *string
	options []string
}

func newEnumFlag(target *string, options ...string) *enumValue {
	return &enumValue{
		value:   target,
		options: options,
	}
}

func (a *enumValue) String() string {
	return *a.value
}

func (a *enumValue) Set(value string) error {
	for _, v := range a.options {
		if v == value {
			*a.value = value
			return nil
		}
	}
	return fmt.Errorf("enum value must be one of %s, got '%s'", strings.Join(a.options, ","), value)
}

func (e *enumValue) Get() interface{} {
	return (string)(*e.value)
}

// -- []string Enum Value
type enumsValue struct {
	value   *[]string
	options []string
}

func newEnumsFlag(target *[]string, options ...string) *enumsValue {
	return &enumsValue{
		value:   target,
		options: options,
	}
}

func (s *enumsValue) Set(value string) error {
	for _, v := range s.options {
		if v == value {
			*s.value = append(*s.value, value)
			return nil
		}
	}
	return fmt.Errorf("enum value must be one of %s, got '%s'", strings.Join(s.options, ","), value)
}

func (e *enumsValue) Get() interface{} {
	return ([]string)(*e.value)
}

func (s *enumsValue) String() string {
	return strings.Join(*s.value, ",")
}

func (s *enumsValue) IsCumulative() bool {
	return true
}

// -- units.Base2Bytes Value
type bytesValue units.Base2Bytes

func newBytesValue(p *units.Base2Bytes) *bytesValue {
	return (*bytesValue)(p)
}

func (d *bytesValue) Set(s string) error {
	v, err := units.ParseBase2Bytes(s)
	*d = bytesValue(v)
	return err
}

func (d *bytesValue) Get() interface{} { return units.Base2Bytes(*d) }

func (d *bytesValue) String() string { return (*units.Base2Bytes)(d).String() }

func newExistingFileValue(target *string) *fileStatValue {
	return newFileStatValue(target, func(s os.FileInfo) error {
		if s.IsDir() {
			return fmt.Errorf("'%s' is a directory", s.Name())
		}
		return nil
	})
}

func newExistingDirValue(target *string) *fileStatValue {
	return newFileStatValue(target, func(s os.FileInfo) error {
		if !s.IsDir() {
			return fmt.Errorf("'%s' is a file", s.Name())
		}
		return nil
	})
}

func newExistingFileOrDirValue(target *string) *fileStatValue {
	return newFileStatValue(target, func(s os.FileInfo) error { return nil })
}

type counterValue int

func newCounterValue(n *int) *counterValue {
	return (*counterValue)(n)
}

func (c *counterValue) Set(s string) error {
	*c++
	return nil
}

func (c *counterValue) Get() interface{} { return (int)(*c) }
func (c *counterValue) IsBoolFlag() bool { return true }
func (c *counterValue) String() string   { return fmt.Sprintf("%d", *c) }
<|endoftext|>


<commit_before>package youtubedl

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
)

// Error youtubedl specific error
type Error string

func (e Error) Error() string {
	return string(e)
}

// Info youtubedl json, thumbnail bytes and raw JSON
type Info struct {
	Artist   string `json:"artist"`
	Uploader string `json:"uploader"`
	Creator  string `json:"creator"`

	Title       string   `json:"title"`
	Description string   `json:"description"`
	Duration    float64  `json:"duration"`
	Thumbnail   string   `json:"thumbnail"`
	Formats     []Format `json:"formats"`

	// not unmarshalled, populated from image thumbnail file
	ThumbnailBytes []byte `json:"-"`

	// private, save raw json to be used later when downloading
	rawJSON []byte
}

// Format youtubedl downloadable format
type Format struct {
	FormatID string  `json:"format_id"`
	Protocol string  `json:"protocol"`
	Ext      string  `json:"ext"`
	ACodec   string  `json:"acodec"`
	VCodec   string  `json:"vcodec"`
	TBR      float64 `json:"tbr"`
	ABR      float64 `json:"abr"`
	VBR      float64 `json:"vbr"`

	NormBR     float64
	NormACodec string
	NormVCodec string
}

func (f Format) String() string {
	return fmt.Sprintf("%s:%s:%s a:%s:%f v:%s:%f %f",
		f.FormatID,
		f.Protocol,
		f.Ext,
		f.NormACodec,
		f.ABR,
		f.NormVCodec,
		f.VBR,
		f.NormBR,
	)
}

// guess codec from fuzzy codec name
func normalizeCodecName(c string) string {
	codecNameNormalizeMap := map[string]string{
		"none": "",
		"avc1": "h264",
		"mp4a": "aac",
		"mp4v": "h264",
		"h265": "hevc",
	}

	// "  NAME.something  " -> "name"
	c = strings.Trim(c, " ")
	c = strings.ToLower(c)
	p := strings.SplitN(c, ".", 2)
	c = p[0]

	if n, ok := codecNameNormalizeMap[c]; ok {
		return n
	}

	return c
}

// guess codecs based on ext
func codecFromExt(ext string) (acodec string, vcodec string) {
	switch strings.ToLower(ext) {
	case "mp3":
		return "mp3", ""
	case "mp4":
		return "aac", "h264"
	case "flv":
		return "aac", "h264"
	default:
		return "", ""
	}
}

func parseInfo(r io.Reader) (info Info, err error) {
	info = Info{}

	info.rawJSON, err = ioutil.ReadAll(r)
	if err != nil {
		return Info{}, err
	}

	if err := json.Unmarshal(info.rawJSON, &info); err != nil {
		return Info{}, err
	}

	for i, _ := range info.Formats {
		f := &info.Formats[i]

		f.NormACodec = normalizeCodecName(f.ACodec)
		f.NormVCodec = normalizeCodecName(f.VCodec)

		extACodec, extVCodec := codecFromExt(f.Ext)
		if f.ACodec == "" {
			f.NormACodec = extACodec
		}
		if f.VCodec == "" {
			f.NormVCodec = extVCodec
		}

		if f.TBR != 0 {
			f.NormBR = f.TBR
		} else {
			f.NormBR = f.ABR + f.VBR
		}
	}

	return info, nil
}

// NewFromURL new Info downloaded from URL using context
func NewFromURL(ctx context.Context, url string, stdout io.Writer) (info Info, err error) {
	tempPath, _ := ioutil.TempDir("", "ydls-youtubedl")
	defer os.RemoveAll(tempPath)

	cmd := exec.CommandContext(
		ctx,
		"youtube-dl",
		"--no-call-home",
		"--no-cache-dir",
		"--skip-download",
		"--write-info-json",
		"--write-thumbnail",
		"--restrict-filenames",
		// don't base output filename on source info
		"--output", "source",
		// provide URL via stdin for security, youtube-dl has some run command args
		"--batch-file", "-",
	)
	cmd.Dir = tempPath
	cmd.Stdout = stdout
	cmdStderr, cmdStderrErr := cmd.StderrPipe()
	if cmdStderrErr != nil {
		return Info{}, cmdStderrErr
	}
	cmdStdin, cmdStdinErr := cmd.StdinPipe()
	if cmdStdinErr != nil {
		return Info{}, cmdStdinErr
	}

	if err := cmd.Start(); err != nil {
		return Info{}, err
	}
	defer cmd.Wait()

	fmt.Fprintln(cmdStdin, url)
	cmdStdin.Close()

	stderrLineScanner := bufio.NewScanner(cmdStderr)
	for stderrLineScanner.Scan() {
		const errorPrefix = "ERROR: "
		line := stderrLineScanner.Text()
		if strings.HasPrefix(line, errorPrefix) {
			return Info{}, Error(line[len(errorPrefix):])
		}
	}

	return NewFromPath(tempPath)
}

// NewFromPath new Info from path with JSON and optional image
func NewFromPath(infoPath string) (info Info, err error) {
	files, err := ioutil.ReadDir(infoPath)
	if err != nil {
		return Info{}, err
	}

	infoJSONPath := ""
	thumbnailPath := ""
	for _, f := range files {
		ext := path.Ext(f.Name())
		if ext == ".json" {
			infoJSONPath = path.Join(infoPath, f.Name())
		} else if ext == ".png" || ext == ".jpg" || ext == ".jpeg" {
			thumbnailPath = path.Join(infoPath, f.Name())
		}
	}

	if infoPath == "" {
		return Info{}, fmt.Errorf("no info json found")
	}

	file, err := os.Open(infoJSONPath)
	if err != nil {
		return Info{}, err
	}
	defer file.Close()
	info, err = parseInfo(file)
	if err != nil {
		return Info{}, err
	}

	if thumbnailPath != "" {
		info.ThumbnailBytes, err = ioutil.ReadFile(thumbnailPath)
		if err != nil {
			return Info{}, fmt.Errorf("failed to read thumbnail")
		}
	}

	return info, nil
}

// DownloadResult download result
type DownloadResult struct {
	Reader io.ReadCloser
	waitCh chan struct{}
}

// Download format matched by filter
func (info Info) Download(ctx context.Context, filter string, stderr io.Writer) (*DownloadResult, error) {
	tempPath, tempErr := ioutil.TempDir("", "ydls-youtubedl")
	if tempErr != nil {
		return nil, tempErr
	}
	jsonTempPath := path.Join(tempPath, "info.json")
	if err := ioutil.WriteFile(jsonTempPath, info.rawJSON, 0644); err != nil {
		os.RemoveAll(tempPath)
		return nil, err
	}

	dr := &DownloadResult{
		waitCh: make(chan struct{}),
	}

	cmd := exec.CommandContext(
		ctx,
		"youtube-dl",
		"--no-call-home",
		"--no-cache-dir",
		"--restrict-filenames",
		"--load-info", jsonTempPath,
		"-f", filter,
		"-o", "-",
	)
	cmd.Dir = tempPath
	var w io.WriteCloser
	dr.Reader, w = io.Pipe()
	cmd.Stdout = w
	cmd.Stderr = stderr

	if err := cmd.Start(); err != nil {
		return nil, err
	}

	go func() {
		cmd.Wait()
		w.Close()
		os.RemoveAll(tempPath)
		close(dr.waitCh)
	}()

	return dr, nil
}

// Wait for resource cleanup
func (dr *DownloadResult) Wait() {
	<-dr.waitCh
}
<commit_msg>Cleanup and fix vet warning<commit_after>package youtubedl

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
)

// Error youtubedl specific error
type Error string

func (e Error) Error() string {
	return string(e)
}

// Info youtubedl json, thumbnail bytes and raw JSON
type Info struct {
	Artist   string `json:"artist"`
	Uploader string `json:"uploader"`
	Creator  string `json:"creator"`

	Title       string   `json:"title"`
	Description string   `json:"description"`
	Duration    float64  `json:"duration"`
	Thumbnail   string   `json:"thumbnail"`
	Formats     []Format `json:"formats"`

	// not unmarshalled, populated from image thumbnail file
	ThumbnailBytes []byte `json:"-"`

	// private, save raw json to be used later when downloading
	rawJSON []byte
}

// Format youtubedl downloadable format
type Format struct {
	FormatID string  `json:"format_id"`
	Protocol string  `json:"protocol"`
	Ext      string  `json:"ext"`
	ACodec   string  `json:"acodec"`
	VCodec   string  `json:"vcodec"`
	TBR      float64 `json:"tbr"`
	ABR      float64 `json:"abr"`
	VBR      float64 `json:"vbr"`

	NormBR     float64
	NormACodec string
	NormVCodec string
}

func (f Format) String() string {
	return fmt.Sprintf("%s:%s:%s a:%s:%f v:%s:%f %f",
		f.FormatID,
		f.Protocol,
		f.Ext,
		f.NormACodec,
		f.ABR,
		f.NormVCodec,
		f.VBR,
		f.NormBR,
	)
}

// guess codec from fuzzy codec name
func normalizeCodecName(c string) string {
	codecNameNormalizeMap := map[string]string{
		"none": "",
		"avc1": "h264",
		"mp4a": "aac",
		"mp4v": "h264",
		"h265": "hevc",
	}

	// "  NAME.something  " -> "name"
	c = strings.Trim(c, " ")
	c = strings.ToLower(c)
	p := strings.SplitN(c, ".", 2)
	c = p[0]

	if n, ok := codecNameNormalizeMap[c]; ok {
		return n
	}

	return c
}

// guess codecs based on ext
func codecFromExt(ext string) (acodec string, vcodec string) {
	switch strings.ToLower(ext) {
	case "mp3":
		return "mp3", ""
	case "mp4":
		return "aac", "h264"
	case "flv":
		return "aac", "h264"
	default:
		return "", ""
	}
}

func parseInfo(r io.Reader) (info Info, err error) {
	info = Info{}

	info.rawJSON, err = ioutil.ReadAll(r)
	if err != nil {
		return Info{}, err
	}

	if err := json.Unmarshal(info.rawJSON, &info); err != nil {
		return Info{}, err
	}

	for i := range info.Formats {
		f := &info.Formats[i]

		f.NormACodec = normalizeCodecName(f.ACodec)
		f.NormVCodec = normalizeCodecName(f.VCodec)

		extACodec, extVCodec := codecFromExt(f.Ext)
		if f.ACodec == "" {
			f.NormACodec = extACodec
		}
		if f.VCodec == "" {
			f.NormVCodec = extVCodec
		}

		if f.TBR != 0 {
			f.NormBR = f.TBR
		} else {
			f.NormBR = f.ABR + f.VBR
		}
	}

	return info, nil
}

// NewFromURL new Info downloaded from URL using context
func NewFromURL(ctx context.Context, url string, stdout io.Writer) (info Info, err error) {
	tempPath, _ := ioutil.TempDir("", "ydls-youtubedl")
	defer os.RemoveAll(tempPath)

	cmd := exec.CommandContext(
		ctx,
		"youtube-dl",
		"--no-call-home",
		"--no-cache-dir",
		"--skip-download",
		"--write-info-json",
		"--write-thumbnail",
		"--restrict-filenames",
		// don't base output filename on source info
		"--output", "source",
		// provide URL via stdin for security, youtube-dl has some run command args
		"--batch-file", "-",
	)
	cmd.Dir = tempPath
	cmd.Stdout = stdout
	cmdStderr, cmdStderrErr := cmd.StderrPipe()
	if cmdStderrErr != nil {
		return Info{}, cmdStderrErr
	}
	cmdStdin, cmdStdinErr := cmd.StdinPipe()
	if cmdStdinErr != nil {
		return Info{}, cmdStdinErr
	}

	if err := cmd.Start(); err != nil {
		return Info{}, err
	}
	defer cmd.Wait()

	fmt.Fprintln(cmdStdin, url)
	cmdStdin.Close()

	stderrLineScanner := bufio.NewScanner(cmdStderr)
	for stderrLineScanner.Scan() {
		const errorPrefix = "ERROR: "
		line := stderrLineScanner.Text()
		if strings.HasPrefix(line, errorPrefix) {
			return Info{}, Error(line[len(errorPrefix):])
		}
	}

	return NewFromPath(tempPath)
}

// NewFromPath new Info from path with JSON and optional image
func NewFromPath(infoPath string) (info Info, err error) {
	files, err := ioutil.ReadDir(infoPath)
	if err != nil {
		return Info{}, err
	}

	infoJSONPath := ""
	thumbnailPath := ""
	for _, f := range files {
		ext := path.Ext(f.Name())
		if ext == ".json" {
			infoJSONPath = path.Join(infoPath, f.Name())
		} else if ext == ".png" || ext == ".jpg" || ext == ".jpeg" {
			thumbnailPath = path.Join(infoPath, f.Name())
		}
	}

	if infoPath == "" {
		return Info{}, fmt.Errorf("no info json found")
	}

	file, err := os.Open(infoJSONPath)
	if err != nil {
		return Info{}, err
	}
	defer file.Close()
	info, err = parseInfo(file)
	if err != nil {
		return Info{}, err
	}

	if thumbnailPath != "" {
		info.ThumbnailBytes, err = ioutil.ReadFile(thumbnailPath)
		if err != nil {
			return Info{}, fmt.Errorf("failed to read thumbnail")
		}
	}

	return info, nil
}

// DownloadResult download result
type DownloadResult struct {
	Reader io.ReadCloser
	waitCh chan struct{}
}

// Wait for resource cleanup
func (dr *DownloadResult) Wait() {
	<-dr.waitCh
}

// Download format matched by filter
func (info Info) Download(ctx context.Context, filter string, stderr io.Writer) (*DownloadResult, error) {
	tempPath, tempErr := ioutil.TempDir("", "ydls-youtubedl")
	if tempErr != nil {
		return nil, tempErr
	}
	jsonTempPath := path.Join(tempPath, "info.json")
	if err := ioutil.WriteFile(jsonTempPath, info.rawJSON, 0644); err != nil {
		os.RemoveAll(tempPath)
		return nil, err
	}

	dr := &DownloadResult{
		waitCh: make(chan struct{}),
	}

	cmd := exec.CommandContext(
		ctx,
		"youtube-dl",
		"--no-call-home",
		"--no-cache-dir",
		"--restrict-filenames",
		"--load-info", jsonTempPath,
		"-f", filter,
		"-o", "-",
	)
	cmd.Dir = tempPath
	var w io.WriteCloser
	dr.Reader, w = io.Pipe()
	cmd.Stdout = w
	cmd.Stderr = stderr

	if err := cmd.Start(); err != nil {
		return nil, err
	}

	go func() {
		cmd.Wait()
		w.Close()
		os.RemoveAll(tempPath)
		close(dr.waitCh)
	}()

	return dr, nil
}
<|endoftext|>


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

package rsyslog_test

import (
	"encoding/pem"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
	apirsyslog "launchpad.net/juju-core/state/api/rsyslog"
	"launchpad.net/juju-core/state/apiserver/common"
	commontesting "launchpad.net/juju-core/state/apiserver/common/testing"
	"launchpad.net/juju-core/state/apiserver/rsyslog"
	apiservertesting "launchpad.net/juju-core/state/apiserver/testing"
	coretesting "launchpad.net/juju-core/testing"
	jc "launchpad.net/juju-core/testing/checkers"
)

type rsyslogSuite struct {
	testing.JujuConnSuite
	*commontesting.EnvironWatcherTest
	authorizer apiservertesting.FakeAuthorizer
	resources  *common.Resources
}

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

func (s *rsyslogSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)
	s.authorizer = apiservertesting.FakeAuthorizer{
		LoggedIn:       true,
		EnvironManager: true,
	}
	s.resources = common.NewResources()
	api, err := rsyslog.NewRsyslogAPI(s.State, s.resources, s.authorizer)
	c.Assert(err, gc.IsNil)
	s.EnvironWatcherTest = commontesting.NewEnvironWatcherTest(
		api, s.State, s.resources, commontesting.NoSecrets)
}

func verifyRsyslogCACert(c *gc.C, st *apirsyslog.State, expected []byte) {
	cfg, err := st.EnvironConfig()
	c.Assert(err, gc.IsNil)
	c.Assert(cfg.RsyslogCACert(), gc.DeepEquals, expected)
}

func (s *rsyslogSuite) TestSetRsyslogCert(c *gc.C) {
	st, _ := s.OpenAPIAsNewMachine(c, state.JobManageEnviron)
	err := st.Rsyslog().SetRsyslogCert([]byte(coretesting.CACert))
	c.Assert(err, gc.IsNil)
	verifyRsyslogCACert(c, st.Rsyslog(), []byte(coretesting.CACert))
}

func (s *rsyslogSuite) TestSetRsyslogCertNil(c *gc.C) {
	st, _ := s.OpenAPIAsNewMachine(c, state.JobManageEnviron)
	err := st.Rsyslog().SetRsyslogCert(nil)
	c.Assert(err, gc.ErrorMatches, "no certificates found")
	verifyRsyslogCACert(c, st.Rsyslog(), nil)
}

func (s *rsyslogSuite) TestSetRsyslogCertInvalid(c *gc.C) {
	st, _ := s.OpenAPIAsNewMachine(c, state.JobManageEnviron)
	err := st.Rsyslog().SetRsyslogCert(pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: []byte("not a valid certificate"),
	}))
	c.Assert(err, gc.ErrorMatches, "asn1: structure error.*")
	verifyRsyslogCACert(c, st.Rsyslog(), nil)
}

func (s *rsyslogSuite) TestSetRsyslogCertPerms(c *gc.C) {
	st, _ := s.OpenAPIAsNewMachine(c, state.JobHostUnits)
	err := st.Rsyslog().SetRsyslogCert([]byte(coretesting.CACert))
	c.Assert(err, gc.ErrorMatches, "invalid entity name or password")
	c.Assert(err, jc.Satisfies, params.IsCodeUnauthorized)
	// Verify no change was effected.
	verifyRsyslogCACert(c, st.Rsyslog(), nil)
}
<commit_msg>state/apiserver/rsyslog: fix ASN.1 test<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package rsyslog_test

import (
	"encoding/pem"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/api/params"
	apirsyslog "launchpad.net/juju-core/state/api/rsyslog"
	"launchpad.net/juju-core/state/apiserver/common"
	commontesting "launchpad.net/juju-core/state/apiserver/common/testing"
	"launchpad.net/juju-core/state/apiserver/rsyslog"
	apiservertesting "launchpad.net/juju-core/state/apiserver/testing"
	coretesting "launchpad.net/juju-core/testing"
	jc "launchpad.net/juju-core/testing/checkers"
)

type rsyslogSuite struct {
	testing.JujuConnSuite
	*commontesting.EnvironWatcherTest
	authorizer apiservertesting.FakeAuthorizer
	resources  *common.Resources
}

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

func (s *rsyslogSuite) SetUpTest(c *gc.C) {
	s.JujuConnSuite.SetUpTest(c)
	s.authorizer = apiservertesting.FakeAuthorizer{
		LoggedIn:       true,
		EnvironManager: true,
	}
	s.resources = common.NewResources()
	api, err := rsyslog.NewRsyslogAPI(s.State, s.resources, s.authorizer)
	c.Assert(err, gc.IsNil)
	s.EnvironWatcherTest = commontesting.NewEnvironWatcherTest(
		api, s.State, s.resources, commontesting.NoSecrets)
}

func verifyRsyslogCACert(c *gc.C, st *apirsyslog.State, expected []byte) {
	cfg, err := st.EnvironConfig()
	c.Assert(err, gc.IsNil)
	c.Assert(cfg.RsyslogCACert(), gc.DeepEquals, expected)
}

func (s *rsyslogSuite) TestSetRsyslogCert(c *gc.C) {
	st, _ := s.OpenAPIAsNewMachine(c, state.JobManageEnviron)
	err := st.Rsyslog().SetRsyslogCert([]byte(coretesting.CACert))
	c.Assert(err, gc.IsNil)
	verifyRsyslogCACert(c, st.Rsyslog(), []byte(coretesting.CACert))
}

func (s *rsyslogSuite) TestSetRsyslogCertNil(c *gc.C) {
	st, _ := s.OpenAPIAsNewMachine(c, state.JobManageEnviron)
	err := st.Rsyslog().SetRsyslogCert(nil)
	c.Assert(err, gc.ErrorMatches, "no certificates found")
	verifyRsyslogCACert(c, st.Rsyslog(), nil)
}

func (s *rsyslogSuite) TestSetRsyslogCertInvalid(c *gc.C) {
	st, _ := s.OpenAPIAsNewMachine(c, state.JobManageEnviron)
	err := st.Rsyslog().SetRsyslogCert(pem.EncodeToMemory(&pem.Block{
		Type:  "CERTIFICATE",
		Bytes: []byte("not a valid certificate"),
	}))
	c.Assert(err, gc.ErrorMatches, ".*structure error.*")
	verifyRsyslogCACert(c, st.Rsyslog(), nil)
}

func (s *rsyslogSuite) TestSetRsyslogCertPerms(c *gc.C) {
	st, _ := s.OpenAPIAsNewMachine(c, state.JobHostUnits)
	err := st.Rsyslog().SetRsyslogCert([]byte(coretesting.CACert))
	c.Assert(err, gc.ErrorMatches, "invalid entity name or password")
	c.Assert(err, jc.Satisfies, params.IsCodeUnauthorized)
	// Verify no change was effected.
	verifyRsyslogCACert(c, st.Rsyslog(), nil)
}
<|endoftext|>


<commit_before>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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.

// A command that reads all blobs necessary for verifying the directory
// structure rooted at a set of backup root scores, ensuring that the entire
// directory structure is intact in GCS.
//
// Optionally, all file content is also read and verified. This is less
// important than verifying directory connectedness if we trust that GCS does
// not corrupt object metadata (where we store expected CRC32C and MD5) and
// does correctly report the object's CRC32C and MD5 sums in listings,
// verifying them periodically.

package main

import (
	"fmt"
	"log"
	"runtime"
	"strings"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/util"
	"github.com/jacobsa/comeback/verify"
	"github.com/jacobsa/comeback/wiring"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"golang.org/x/net/context"
)

var cmdVerify = &Command{
	Name: "verify",
}

// TODO(jacobsa): Get these automatically from the registry.
var fRoots = cmdVerify.Flags.String(
	"roots",
	"",
	"Comma-separated list of backup root scores to verify.")

var fFast = cmdVerify.Flags.Bool(
	"fast",
	false,
	"When set, don't verify file content.")

func init() {
	cmdVerify.Run = runVerify // Break flag-related dependency loop.
}

type loggingVisitor struct {
	wrapped graph.Visitor
}

func (v *loggingVisitor) Visit(
	ctx context.Context,
	node string) (adjacent []string, err error) {
	// Call through.
	adjacent, err = v.wrapped.Visit(ctx, node)
	if err != nil {
		return
	}

	// Log.
	log.Printf("%s -> %s", node, strings.Join(adjacent, " "))

	return
}

// List blob.ListScores, but returns a slice instead of writing into a channel.
func listAllScores(
	ctx context.Context,
	bucket gcs.Bucket,
	namePrefix string) (scores []blob.Score, err error) {
	b := syncutil.NewBundle(context.Background())
	defer func() { err = b.Join() }()

	// List scores into a channel.
	scoreChan := make(chan blob.Score, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(scoreChan)
		err = blob.ListScores(ctx, bucket, wiring.BlobObjectNamePrefix, scoreChan)
		if err != nil {
			err = fmt.Errorf("ListScores: %v", err)
			return
		}

		return
	})

	// Accumulate into the slice.
	b.Add(func(ctx context.Context) (err error) {
		for score := range scoreChan {
			scores = append(scores, score)
		}

		return
	})

	return
}

func runVerify(args []string) {
	// Allow parallelism.
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Die on error.
	var err error
	defer func() {
		if err != nil {
			log.Fatalln(err)
		}
	}()

	// Read flags.
	readFiles := !*fFast

	if *fRoots == "" {
		err = fmt.Errorf("You must set --roots.")
		return
	}

	rootHexScores := strings.Split(*fRoots, ",")
	var roots []string
	for _, hexScore := range rootHexScores {
		var score blob.Score
		score, err = blob.ParseHexScore(hexScore)
		if err != nil {
			err = fmt.Errorf("Invalid root %q: %v", hexScore, err)
			return
		}

		roots = append(roots, verify.FormatNodeName(true, score))
	}

	// Grab dependencies.
	bucket := getBucket()
	crypter := getCrypter()

	// Create a blob store.
	blobStore, err := wiring.MakeBlobStore(
		bucket,
		crypter,
		util.NewStringSet())

	if err != nil {
		err = fmt.Errorf("MakeBlobStore: %v", err)
		return
	}

	// List all scores in the bucket, verifying the object record metadata in the
	// process.
	log.Println("Listing scores...")
	knownScores, err := listAllScores(
		context.Background(),
		bucket,
		wiring.BlobObjectNamePrefix)

	if err != nil {
		err = fmt.Errorf("listAllScores: %v", err)
		return
	}

	log.Printf("Listed %d scores.", len(knownScores))

	// Create a graph visitor that perform the verification.
	visitor := verify.NewVisitor(
		readFiles,
		knownScores,
		blobStore)

	visitor = &loggingVisitor{
		wrapped: visitor,
	}

	// Traverse starting at the specified roots. Use an "experimentally
	// determined" parallelism, which in theory should depend on bandwidth-delay
	// products but in practice comes down to when the OS gets cranky about open
	// files.
	const parallelism = 128

	err = graph.Traverse(
		context.Background(),
		parallelism,
		roots,
		visitor)

	if err != nil {
		err = fmt.Errorf("Traverse: %v", err)
		return
	}

	return
}
<commit_msg>Added a count in the output when successful.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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.

// A command that reads all blobs necessary for verifying the directory
// structure rooted at a set of backup root scores, ensuring that the entire
// directory structure is intact in GCS.
//
// Optionally, all file content is also read and verified. This is less
// important than verifying directory connectedness if we trust that GCS does
// not corrupt object metadata (where we store expected CRC32C and MD5) and
// does correctly report the object's CRC32C and MD5 sums in listings,
// verifying them periodically.

package main

import (
	"fmt"
	"log"
	"runtime"
	"strings"
	"sync/atomic"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/util"
	"github.com/jacobsa/comeback/verify"
	"github.com/jacobsa/comeback/wiring"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/syncutil"
	"golang.org/x/net/context"
)

var cmdVerify = &Command{
	Name: "verify",
}

// TODO(jacobsa): Get these automatically from the registry.
var fRoots = cmdVerify.Flags.String(
	"roots",
	"",
	"Comma-separated list of backup root scores to verify.")

var fFast = cmdVerify.Flags.Bool(
	"fast",
	false,
	"When set, don't verify file content.")

func init() {
	cmdVerify.Run = runVerify // Break flag-related dependency loop.
}

////////////////////////////////////////////////////////////////////////
// Visitor types
////////////////////////////////////////////////////////////////////////

type loggingVisitor struct {
	wrapped graph.Visitor
}

func (v *loggingVisitor) Visit(
	ctx context.Context,
	node string) (adjacent []string, err error) {
	// Call through.
	adjacent, err = v.wrapped.Visit(ctx, node)
	if err != nil {
		return
	}

	// Log.
	log.Printf("%s -> %s", node, strings.Join(adjacent, " "))

	return
}

type countingVisitor struct {
	count   *uint64
	wrapped graph.Visitor
}

func (v *countingVisitor) Visit(
	ctx context.Context,
	node string) (adjacent []string, err error) {
	atomic.AddUint64(v.count, 1)
	adjacent, err = v.wrapped.Visit(ctx, node)
	return
}

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

// List blob.ListScores, but returns a slice instead of writing into a channel.
func listAllScores(
	ctx context.Context,
	bucket gcs.Bucket,
	namePrefix string) (scores []blob.Score, err error) {
	b := syncutil.NewBundle(context.Background())
	defer func() { err = b.Join() }()

	// List scores into a channel.
	scoreChan := make(chan blob.Score, 100)
	b.Add(func(ctx context.Context) (err error) {
		defer close(scoreChan)
		err = blob.ListScores(ctx, bucket, wiring.BlobObjectNamePrefix, scoreChan)
		if err != nil {
			err = fmt.Errorf("ListScores: %v", err)
			return
		}

		return
	})

	// Accumulate into the slice.
	b.Add(func(ctx context.Context) (err error) {
		for score := range scoreChan {
			scores = append(scores, score)
		}

		return
	})

	return
}

////////////////////////////////////////////////////////////////////////
// Verify
////////////////////////////////////////////////////////////////////////

func runVerify(args []string) {
	// Allow parallelism.
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Die on error.
	var err error
	defer func() {
		if err != nil {
			log.Fatalln(err)
		}
	}()

	// Read flags.
	readFiles := !*fFast

	if *fRoots == "" {
		err = fmt.Errorf("You must set --roots.")
		return
	}

	rootHexScores := strings.Split(*fRoots, ",")
	var roots []string
	for _, hexScore := range rootHexScores {
		var score blob.Score
		score, err = blob.ParseHexScore(hexScore)
		if err != nil {
			err = fmt.Errorf("Invalid root %q: %v", hexScore, err)
			return
		}

		roots = append(roots, verify.FormatNodeName(true, score))
	}

	// Grab dependencies.
	bucket := getBucket()
	crypter := getCrypter()

	// Create a blob store.
	blobStore, err := wiring.MakeBlobStore(
		bucket,
		crypter,
		util.NewStringSet())

	if err != nil {
		err = fmt.Errorf("MakeBlobStore: %v", err)
		return
	}

	// List all scores in the bucket, verifying the object record metadata in the
	// process.
	log.Println("Listing scores...")
	knownScores, err := listAllScores(
		context.Background(),
		bucket,
		wiring.BlobObjectNamePrefix)

	if err != nil {
		err = fmt.Errorf("listAllScores: %v", err)
		return
	}

	log.Printf("Listed %d scores.", len(knownScores))

	// Create a graph visitor that perform the verification.
	visitor := verify.NewVisitor(
		readFiles,
		knownScores,
		blobStore)

	visitor = &loggingVisitor{
		wrapped: visitor,
	}

	var nodesVisited uint64
	visitor = &countingVisitor{
		count:   &nodesVisited,
		wrapped: visitor,
	}

	// Traverse starting at the specified roots. Use an "experimentally
	// determined" parallelism, which in theory should depend on bandwidth-delay
	// products but in practice comes down to when the OS gets cranky about open
	// files.
	const parallelism = 128

	err = graph.Traverse(
		context.Background(),
		parallelism,
		roots,
		visitor)

	if err != nil {
		err = fmt.Errorf("Traverse: %v", err)
		return
	}

	log.Printf("Successfully visited %v nodes.", nodesVisited)

	return
}
<|endoftext|>


<commit_before>package replicaset

import (
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	// MaxPeers defines the maximum number of peers that mongo supports.
	MaxPeers = 7

	// maxInitiateAttempts is the maximum number of times to attempt
	// replSetInitiate for each call to Initiate.
	maxInitiateAttempts = 10

	// initiateAttemptDelay is the amount of time to sleep between failed
	// attempts to replSetInitiate.
	initiateAttemptDelay = 100 * time.Millisecond

	// rsMembersUnreachableError is the error message returned from mongo
	// when it thinks that replicaset members are unreachable. This can
	// occur if replSetInitiate is executed shortly after starting up mongo.
	rsMembersUnreachableError = "all members and seeds must be reachable to initiate set"
)

var logger = loggo.GetLogger("juju.replicaset")

// Initiate sets up a replica set with the given replica set name with the
// single given member.  It need be called only once for a given mongo replica
// set.  The tags specified will be added as tags on the member that is created
// in the replica set.
//
// Note that you must set DialWithInfo and set Direct = true when dialing into a
// specific non-initiated mongo server.
//
// See http://docs.mongodb.org/manual/reference/method/rs.initiate/ for more
// details.
func Initiate(session *mgo.Session, address, name string, tags map[string]string) error {
	monotonicSession := session.Clone()
	monotonicSession.SetMode(mgo.Monotonic, true)
	cfg := Config{
		Name:    name,
		Version: 1,
		Members: []Member{{
			Id:      1,
			Address: address,
			Tags:    tags,
		}},
	}
	logger.Infof("Initiating replicaset with config %#v", cfg)
	var err error
	for i := 0; i < maxInitiateAttempts; i++ {
		err = monotonicSession.Run(bson.D{{"replSetInitiate", cfg}}, nil)
		if err != nil && err.Error() == rsMembersUnreachableError {
			time.Sleep(initiateAttemptDelay)
			continue
		}
		break
	}
	return err
}

// Member holds configuration information for a replica set member.
//
// See http://docs.mongodb.org/manual/reference/replica-configuration/
// for more details
type Member struct {
	// Id is a unique id for a member in a set.
	Id int `bson:"_id"`

	// Address holds the network address of the member,
	// in the form hostname:port.
	Address string `bson:"host"`

	// Arbiter holds whether the member is an arbiter only.
	// This value is optional; it defaults to false.
	Arbiter *bool `bson:"arbiterOnly,omitempty"`

	// BuildIndexes determines whether the mongod builds indexes on this member.
	// This value is optional; it defaults to true.
	BuildIndexes *bool `bson:"buildIndexes,omitempty"`

	// Hidden determines whether the replica set hides this member from
	// the output of IsMaster.
	// This value is optional; it defaults to false.
	Hidden *bool `bson:"hidden,omitempty"`

	// Priority determines eligibility of a member to become primary.
	// This value is optional; it defaults to 1.
	Priority *float64 `bson:"priority,omitempty"`

	// Tags store additional information about a replica member, often used for
	// customizing read preferences and write concern.
	Tags map[string]string `bson:"tags,omitempty"`

	// SlaveDelay describes the number of seconds behind the master that this
	// replica set member should lag rounded up to the nearest second.
	// This value is optional; it defaults to 0.
	SlaveDelay *time.Duration `bson:"slaveDelay,omitempty"`

	// Votes controls the number of votes a server has in a replica set election.
	// This value is optional; it defaults to 1.
	Votes *int `bson:"votes,omitempty"`
}

func fmtConfigForLog(config *Config) string {
	memberInfo := make([]string, len(config.Members))
	for i, member := range config.Members {
		memberInfo[i] = fmt.Sprintf("Member{%d %q %v}", member.Id, member.Address, member.Tags)

	}
	return fmt.Sprintf("{Name: %s, Version: %d, Members: {%s}}",
		config.Name, config.Version, strings.Join(memberInfo, ", "))
}

// applyRelSetConfig applies the new config to the mongo session. It also logs
// what the changes are. It checks if the replica set changes cause the DB
// connection to be dropped. If so, it Refreshes the session and tries to Ping
// again.
func applyRelSetConfig(cmd string, session *mgo.Session, oldconfig, newconfig *Config) error {
	logger.Debugf("%s() changing replica set\nfrom %s\n  to %s",
		cmd, fmtConfigForLog(oldconfig), fmtConfigForLog(newconfig))
	err := session.Run(bson.D{{"replSetReconfig", newconfig}}, nil)
	// We will only try to Ping 2 times
	for i := 0; i < 2; i++ {
		if err == io.EOF {
			// If the primary changes due to replSetReconfig, then all
			// current connections are dropped.
			// Refreshing should fix us up.
			logger.Debugf("got EOF while running %s(), calling session.Refresh()", cmd)
			session.Refresh()
		} else if err != nil {
			// For all errors that aren't EOF, return immediately
			return err
		}
		// err is either nil or EOF and we called Refresh, so Ping to
		// make sure we're actually connected
		err = session.Ping()
		// Change the command because it is the new command we ran
		cmd = "Ping"
	}
	return err
}

// Add adds the given members to the session's replica set.  Duplicates of
// existing replicas will be ignored.
//
// Members will have their Ids set automatically if they are not already > 0
func Add(session *mgo.Session, members ...Member) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}

	oldconfig := *config
	config.Version++
	max := 0
	for _, member := range config.Members {
		if member.Id > max {
			max = member.Id
		}
	}

outerLoop:
	for _, newMember := range members {
		for _, member := range config.Members {
			if member.Address == newMember.Address {
				// already exists, skip it
				continue outerLoop
			}
		}
		// let the caller specify an id if they want, treat zero as unspecified
		if newMember.Id < 1 {
			max++
			newMember.Id = max
		}
		config.Members = append(config.Members, newMember)
	}
	return applyRelSetConfig("Add", session, &oldconfig, config)
}

// Remove removes members with the given addresses from the replica set. It is
// not an error to remove addresses of non-existent replica set members.
func Remove(session *mgo.Session, addrs ...string) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}
	oldconfig := *config
	config.Version++
	for _, rem := range addrs {
		for n, repl := range config.Members {
			if repl.Address == rem {
				config.Members = append(config.Members[:n], config.Members[n+1:]...)
				break
			}
		}
	}
	return applyRelSetConfig("Remove", session, &oldconfig, config)
}

// Set changes the current set of replica set members.  Members will have their
// ids set automatically if their ids are not already > 0.
func Set(session *mgo.Session, members []Member) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}

	// Copy the current configuration for logging
	oldconfig := *config
	config.Version++

	// Assign ids to members that did not previously exist, starting above the
	// value of the highest id that already existed
	ids := map[string]int{}
	max := 0
	for _, m := range config.Members {
		ids[m.Address] = m.Id
		if m.Id > max {
			max = m.Id
		}
	}

	for x, m := range members {
		if id, ok := ids[m.Address]; ok {
			m.Id = id
		} else if m.Id < 1 {
			max++
			m.Id = max
		}
		members[x] = m
	}

	config.Members = members

	return applyRelSetConfig("Set", session, &oldconfig, config)
}

// Config reports information about the configuration of a given mongo node
type IsMasterResults struct {
	// The following fields hold information about the specific mongodb node.
	IsMaster  bool      `bson:"ismaster"`
	Secondary bool      `bson:"secondary"`
	Arbiter   bool      `bson:"arbiterOnly"`
	Address   string    `bson:"me"`
	LocalTime time.Time `bson:"localTime"`

	// The following fields hold information about the replica set.
	ReplicaSetName string   `bson:"setName"`
	Addresses      []string `bson:"hosts"`
	Arbiters       []string `bson:"arbiters"`
	PrimaryAddress string   `bson:"primary"`
}

// IsMaster returns information about the configuration of the node that
// the given session is connected to.
func IsMaster(session *mgo.Session) (*IsMasterResults, error) {
	results := &IsMasterResults{}
	err := session.Run("isMaster", results)
	if err != nil {
		return nil, err
	}
	return results, nil
}

var ErrMasterNotConfigured = fmt.Errorf("mongo master not configured")

// MasterHostPort returns the "address:port" string for the primary
// mongo server in the replicaset. It returns ErrMasterNotConfigured if
// the replica set has not yet been initiated.
func MasterHostPort(session *mgo.Session) (string, error) {
	results, err := IsMaster(session)
	if err != nil {
		return "", err
	}
	if results.PrimaryAddress == "" {
		return "", ErrMasterNotConfigured
	}
	return results.PrimaryAddress, nil
}

// CurrentMembers returns the current members of the replica set.
func CurrentMembers(session *mgo.Session) ([]Member, error) {
	cfg, err := CurrentConfig(session)
	if err != nil {
		return nil, err
	}
	return cfg.Members, nil
}

// CurrentConfig returns the Config for the given session's replica set.  If
// there is no current config, the error returned will be mgo.ErrNotFound.
func CurrentConfig(session *mgo.Session) (*Config, error) {
	cfg := &Config{}
	monotonicSession := session.Clone()
	monotonicSession.SetMode(mgo.Monotonic, true)
	err := monotonicSession.DB("local").C("system.replset").Find(nil).One(cfg)
	if err == mgo.ErrNotFound {
		return nil, err
	}
	if err != nil {
		return nil, fmt.Errorf("cannot get replset config: %s", err.Error())
	}
	return cfg, nil
}

// Config is the document stored in mongodb that defines the servers in the
// replica set
type Config struct {
	Name    string   `bson:"_id"`
	Version int      `bson:"version"`
	Members []Member `bson:"members"`
}

// CurrentStatus returns the status of the replica set for the given session.
func CurrentStatus(session *mgo.Session) (*Status, error) {
	status := &Status{}
	err := session.Run("replSetGetStatus", status)
	if err != nil {
		return nil, fmt.Errorf("cannot get replica set status: %v", err)
	}
	return status, nil
}

// Status holds data about the status of members of the replica set returned
// from replSetGetStatus
//
// See http://docs.mongodb.org/manual/reference/command/replSetGetStatus/#dbcmd.replSetGetStatus
type Status struct {
	Name    string         `bson:"set"`
	Members []MemberStatus `bson:"members"`
}

// Status holds the status of a replica set member returned from
// replSetGetStatus.
type MemberStatus struct {
	// Id holds the replica set id of the member that the status is describing.
	Id int `bson:"_id"`

	// Address holds address of the member that the status is describing.
	Address string `bson:"name"`

	// Self holds whether this is the status for the member that
	// the session is connected to.
	Self bool `bson:"self"`

	// ErrMsg holds the most recent error or status message received
	// from the member.
	ErrMsg string `bson:"errmsg"`

	// Healthy reports whether the member is up. It is true for the
	// member that the request was made to.
	Healthy bool `bson:"health"`

	// State describes the current state of the member.
	State MemberState `bson:"state"`

	// Uptime describes how long the member has been online.
	Uptime time.Duration `bson:"uptime"`

	// Ping describes the length of time a round-trip packet takes to travel
	// between the remote member and the local instance.  It is zero for the
	// member that the session is connected to.
	Ping time.Duration `bson:"pingMS"`
}

// MemberState represents the state of a replica set member.
// See http://docs.mongodb.org/manual/reference/replica-states/
type MemberState int

const (
	StartupState = iota
	PrimaryState
	SecondaryState
	RecoveringState
	FatalState
	Startup2State
	UnknownState
	ArbiterState
	DownState
	RollbackState
	ShunnedState
)

var memberStateStrings = []string{
	StartupState:    "STARTUP",
	PrimaryState:    "PRIMARY",
	SecondaryState:  "SECONDARY",
	RecoveringState: "RECOVERING",
	FatalState:      "FATAL",
	Startup2State:   "STARTUP2",
	UnknownState:    "UNKNOWN",
	ArbiterState:    "ARBITER",
	DownState:       "DOWN",
	RollbackState:   "ROLLBACK",
	ShunnedState:    "SHUNNED",
}

// String returns a string describing the state.
func (state MemberState) String() string {
	if state < 0 || int(state) >= len(memberStateStrings) {
		return "INVALID_MEMBER_STATE"
	}
	return memberStateStrings[state]
}
<commit_msg>[r=axwalk] replicaset: close cloned sessions<commit_after>package replicaset

import (
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/juju/loggo"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	// MaxPeers defines the maximum number of peers that mongo supports.
	MaxPeers = 7

	// maxInitiateAttempts is the maximum number of times to attempt
	// replSetInitiate for each call to Initiate.
	maxInitiateAttempts = 10

	// initiateAttemptDelay is the amount of time to sleep between failed
	// attempts to replSetInitiate.
	initiateAttemptDelay = 100 * time.Millisecond

	// rsMembersUnreachableError is the error message returned from mongo
	// when it thinks that replicaset members are unreachable. This can
	// occur if replSetInitiate is executed shortly after starting up mongo.
	rsMembersUnreachableError = "all members and seeds must be reachable to initiate set"
)

var logger = loggo.GetLogger("juju.replicaset")

// Initiate sets up a replica set with the given replica set name with the
// single given member.  It need be called only once for a given mongo replica
// set.  The tags specified will be added as tags on the member that is created
// in the replica set.
//
// Note that you must set DialWithInfo and set Direct = true when dialing into a
// specific non-initiated mongo server.
//
// See http://docs.mongodb.org/manual/reference/method/rs.initiate/ for more
// details.
func Initiate(session *mgo.Session, address, name string, tags map[string]string) error {
	monotonicSession := session.Clone()
	defer monotonicSession.Close()
	monotonicSession.SetMode(mgo.Monotonic, true)
	cfg := Config{
		Name:    name,
		Version: 1,
		Members: []Member{{
			Id:      1,
			Address: address,
			Tags:    tags,
		}},
	}
	logger.Infof("Initiating replicaset with config %#v", cfg)
	var err error
	for i := 0; i < maxInitiateAttempts; i++ {
		err = monotonicSession.Run(bson.D{{"replSetInitiate", cfg}}, nil)
		if err != nil && err.Error() == rsMembersUnreachableError {
			time.Sleep(initiateAttemptDelay)
			continue
		}
		break
	}
	return err
}

// Member holds configuration information for a replica set member.
//
// See http://docs.mongodb.org/manual/reference/replica-configuration/
// for more details
type Member struct {
	// Id is a unique id for a member in a set.
	Id int `bson:"_id"`

	// Address holds the network address of the member,
	// in the form hostname:port.
	Address string `bson:"host"`

	// Arbiter holds whether the member is an arbiter only.
	// This value is optional; it defaults to false.
	Arbiter *bool `bson:"arbiterOnly,omitempty"`

	// BuildIndexes determines whether the mongod builds indexes on this member.
	// This value is optional; it defaults to true.
	BuildIndexes *bool `bson:"buildIndexes,omitempty"`

	// Hidden determines whether the replica set hides this member from
	// the output of IsMaster.
	// This value is optional; it defaults to false.
	Hidden *bool `bson:"hidden,omitempty"`

	// Priority determines eligibility of a member to become primary.
	// This value is optional; it defaults to 1.
	Priority *float64 `bson:"priority,omitempty"`

	// Tags store additional information about a replica member, often used for
	// customizing read preferences and write concern.
	Tags map[string]string `bson:"tags,omitempty"`

	// SlaveDelay describes the number of seconds behind the master that this
	// replica set member should lag rounded up to the nearest second.
	// This value is optional; it defaults to 0.
	SlaveDelay *time.Duration `bson:"slaveDelay,omitempty"`

	// Votes controls the number of votes a server has in a replica set election.
	// This value is optional; it defaults to 1.
	Votes *int `bson:"votes,omitempty"`
}

func fmtConfigForLog(config *Config) string {
	memberInfo := make([]string, len(config.Members))
	for i, member := range config.Members {
		memberInfo[i] = fmt.Sprintf("Member{%d %q %v}", member.Id, member.Address, member.Tags)

	}
	return fmt.Sprintf("{Name: %s, Version: %d, Members: {%s}}",
		config.Name, config.Version, strings.Join(memberInfo, ", "))
}

// applyRelSetConfig applies the new config to the mongo session. It also logs
// what the changes are. It checks if the replica set changes cause the DB
// connection to be dropped. If so, it Refreshes the session and tries to Ping
// again.
func applyRelSetConfig(cmd string, session *mgo.Session, oldconfig, newconfig *Config) error {
	logger.Debugf("%s() changing replica set\nfrom %s\n  to %s",
		cmd, fmtConfigForLog(oldconfig), fmtConfigForLog(newconfig))
	err := session.Run(bson.D{{"replSetReconfig", newconfig}}, nil)
	// We will only try to Ping 2 times
	for i := 0; i < 2; i++ {
		if err == io.EOF {
			// If the primary changes due to replSetReconfig, then all
			// current connections are dropped.
			// Refreshing should fix us up.
			logger.Debugf("got EOF while running %s(), calling session.Refresh()", cmd)
			session.Refresh()
		} else if err != nil {
			// For all errors that aren't EOF, return immediately
			return err
		}
		// err is either nil or EOF and we called Refresh, so Ping to
		// make sure we're actually connected
		err = session.Ping()
		// Change the command because it is the new command we ran
		cmd = "Ping"
	}
	return err
}

// Add adds the given members to the session's replica set.  Duplicates of
// existing replicas will be ignored.
//
// Members will have their Ids set automatically if they are not already > 0
func Add(session *mgo.Session, members ...Member) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}

	oldconfig := *config
	config.Version++
	max := 0
	for _, member := range config.Members {
		if member.Id > max {
			max = member.Id
		}
	}

outerLoop:
	for _, newMember := range members {
		for _, member := range config.Members {
			if member.Address == newMember.Address {
				// already exists, skip it
				continue outerLoop
			}
		}
		// let the caller specify an id if they want, treat zero as unspecified
		if newMember.Id < 1 {
			max++
			newMember.Id = max
		}
		config.Members = append(config.Members, newMember)
	}
	return applyRelSetConfig("Add", session, &oldconfig, config)
}

// Remove removes members with the given addresses from the replica set. It is
// not an error to remove addresses of non-existent replica set members.
func Remove(session *mgo.Session, addrs ...string) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}
	oldconfig := *config
	config.Version++
	for _, rem := range addrs {
		for n, repl := range config.Members {
			if repl.Address == rem {
				config.Members = append(config.Members[:n], config.Members[n+1:]...)
				break
			}
		}
	}
	return applyRelSetConfig("Remove", session, &oldconfig, config)
}

// Set changes the current set of replica set members.  Members will have their
// ids set automatically if their ids are not already > 0.
func Set(session *mgo.Session, members []Member) error {
	config, err := CurrentConfig(session)
	if err != nil {
		return err
	}

	// Copy the current configuration for logging
	oldconfig := *config
	config.Version++

	// Assign ids to members that did not previously exist, starting above the
	// value of the highest id that already existed
	ids := map[string]int{}
	max := 0
	for _, m := range config.Members {
		ids[m.Address] = m.Id
		if m.Id > max {
			max = m.Id
		}
	}

	for x, m := range members {
		if id, ok := ids[m.Address]; ok {
			m.Id = id
		} else if m.Id < 1 {
			max++
			m.Id = max
		}
		members[x] = m
	}

	config.Members = members

	return applyRelSetConfig("Set", session, &oldconfig, config)
}

// Config reports information about the configuration of a given mongo node
type IsMasterResults struct {
	// The following fields hold information about the specific mongodb node.
	IsMaster  bool      `bson:"ismaster"`
	Secondary bool      `bson:"secondary"`
	Arbiter   bool      `bson:"arbiterOnly"`
	Address   string    `bson:"me"`
	LocalTime time.Time `bson:"localTime"`

	// The following fields hold information about the replica set.
	ReplicaSetName string   `bson:"setName"`
	Addresses      []string `bson:"hosts"`
	Arbiters       []string `bson:"arbiters"`
	PrimaryAddress string   `bson:"primary"`
}

// IsMaster returns information about the configuration of the node that
// the given session is connected to.
func IsMaster(session *mgo.Session) (*IsMasterResults, error) {
	results := &IsMasterResults{}
	err := session.Run("isMaster", results)
	if err != nil {
		return nil, err
	}
	return results, nil
}

var ErrMasterNotConfigured = fmt.Errorf("mongo master not configured")

// MasterHostPort returns the "address:port" string for the primary
// mongo server in the replicaset. It returns ErrMasterNotConfigured if
// the replica set has not yet been initiated.
func MasterHostPort(session *mgo.Session) (string, error) {
	results, err := IsMaster(session)
	if err != nil {
		return "", err
	}
	if results.PrimaryAddress == "" {
		return "", ErrMasterNotConfigured
	}
	return results.PrimaryAddress, nil
}

// CurrentMembers returns the current members of the replica set.
func CurrentMembers(session *mgo.Session) ([]Member, error) {
	cfg, err := CurrentConfig(session)
	if err != nil {
		return nil, err
	}
	return cfg.Members, nil
}

// CurrentConfig returns the Config for the given session's replica set.  If
// there is no current config, the error returned will be mgo.ErrNotFound.
func CurrentConfig(session *mgo.Session) (*Config, error) {
	cfg := &Config{}
	monotonicSession := session.Clone()
	defer monotonicSession.Close()
	monotonicSession.SetMode(mgo.Monotonic, true)
	err := monotonicSession.DB("local").C("system.replset").Find(nil).One(cfg)
	if err == mgo.ErrNotFound {
		return nil, err
	}
	if err != nil {
		return nil, fmt.Errorf("cannot get replset config: %s", err.Error())
	}
	return cfg, nil
}

// Config is the document stored in mongodb that defines the servers in the
// replica set
type Config struct {
	Name    string   `bson:"_id"`
	Version int      `bson:"version"`
	Members []Member `bson:"members"`
}

// CurrentStatus returns the status of the replica set for the given session.
func CurrentStatus(session *mgo.Session) (*Status, error) {
	status := &Status{}
	err := session.Run("replSetGetStatus", status)
	if err != nil {
		return nil, fmt.Errorf("cannot get replica set status: %v", err)
	}
	return status, nil
}

// Status holds data about the status of members of the replica set returned
// from replSetGetStatus
//
// See http://docs.mongodb.org/manual/reference/command/replSetGetStatus/#dbcmd.replSetGetStatus
type Status struct {
	Name    string         `bson:"set"`
	Members []MemberStatus `bson:"members"`
}

// Status holds the status of a replica set member returned from
// replSetGetStatus.
type MemberStatus struct {
	// Id holds the replica set id of the member that the status is describing.
	Id int `bson:"_id"`

	// Address holds address of the member that the status is describing.
	Address string `bson:"name"`

	// Self holds whether this is the status for the member that
	// the session is connected to.
	Self bool `bson:"self"`

	// ErrMsg holds the most recent error or status message received
	// from the member.
	ErrMsg string `bson:"errmsg"`

	// Healthy reports whether the member is up. It is true for the
	// member that the request was made to.
	Healthy bool `bson:"health"`

	// State describes the current state of the member.
	State MemberState `bson:"state"`

	// Uptime describes how long the member has been online.
	Uptime time.Duration `bson:"uptime"`

	// Ping describes the length of time a round-trip packet takes to travel
	// between the remote member and the local instance.  It is zero for the
	// member that the session is connected to.
	Ping time.Duration `bson:"pingMS"`
}

// MemberState represents the state of a replica set member.
// See http://docs.mongodb.org/manual/reference/replica-states/
type MemberState int

const (
	StartupState = iota
	PrimaryState
	SecondaryState
	RecoveringState
	FatalState
	Startup2State
	UnknownState
	ArbiterState
	DownState
	RollbackState
	ShunnedState
)

var memberStateStrings = []string{
	StartupState:    "STARTUP",
	PrimaryState:    "PRIMARY",
	SecondaryState:  "SECONDARY",
	RecoveringState: "RECOVERING",
	FatalState:      "FATAL",
	Startup2State:   "STARTUP2",
	UnknownState:    "UNKNOWN",
	ArbiterState:    "ARBITER",
	DownState:       "DOWN",
	RollbackState:   "ROLLBACK",
	ShunnedState:    "SHUNNED",
}

// String returns a string describing the state.
func (state MemberState) String() string {
	if state < 0 || int(state) >= len(memberStateStrings) {
		return "INVALID_MEMBER_STATE"
	}
	return memberStateStrings[state]
}
<|endoftext|>