<commit_before>package dbtest

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"os/exec"
	"strconv"
	"time"

	"gopkg.in/mgo.v2-unstable"
	"gopkg.in/tomb.v2"
)

// DBServer controls a MongoDB server process to be used within test suites.
//
// The test server is started when Session is called the first time and should
// remain running for the duration of all tests, with the Wipe method being
// called between tests (before each of them) to clear stored data. After all tests
// are done, the Stop method should be called to stop the test server.
//
// Before the DBServer is used the SetPath method must be called to define
// the location for the database files to be stored.
type DBServer struct {
	session *mgo.Session
	output  bytes.Buffer
	server  *exec.Cmd
	dbpath  string
	host    string
	tomb    tomb.Tomb
}

// SetPath defines the path to the directory where the database files will be
// stored if it is started. The directory path itself is not created or removed
// by the test helper.
func (dbs *DBServer) SetPath(dbpath string) {
	dbs.dbpath = dbpath
}

func (dbs *DBServer) start() {
	if dbs.server != nil {
		panic("DBServer already started")
	}
	if dbs.dbpath == "" {
		panic("DBServer.SetPath must be called before using the server")
	}
	mgo.SetStats(true)
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic("unable to listen on a local address: " + err.Error())
	}
	addr := l.Addr().(*net.TCPAddr)
	l.Close()
	dbs.host = addr.String()

	args := []string{
		"--dbpath", dbs.dbpath,
		"--bind_ip", "127.0.0.1",
		"--port", strconv.Itoa(addr.Port),
		"--nssize", "1",
		"--noprealloc",
		"--smallfiles",
		"--nojournal",
	}
	dbs.tomb = tomb.Tomb{}
	dbs.server = exec.Command("mongod", args...)
	dbs.server.Stdout = &dbs.output
	dbs.server.Stderr = &dbs.output
	err = dbs.server.Start()
	if err != nil {
		panic(err)
	}
	dbs.tomb.Go(dbs.monitor)
	dbs.Wipe()
}

func (dbs *DBServer) monitor() error {
	dbs.server.Process.Wait()
	if dbs.tomb.Alive() {
		// Present some debugging information.
		fmt.Fprintf(os.Stderr, "---- mongod process died unexpectedly:\n")
		fmt.Fprintf(os.Stderr, "%s", dbs.output.Bytes())
		fmt.Fprintf(os.Stderr, "---- mongod processes running right now:\n")
		cmd := exec.Command("/bin/sh", "-c", "ps auxw | grep mongod")
		cmd.Stdout = os.Stderr
		cmd.Stderr = os.Stderr
		cmd.Run()
		fmt.Fprintf(os.Stderr, "----------------------------------------\n")

		panic("mongod process died unexpectedly")
	}
	return nil
}

// Stop stops the test server process, if it is running.
//
// It's okay to call Stop multiple times. After the test server is
// stopped it cannot be restarted.
//
// All database sessions must be closed before or while the Stop method
// is running. Otherwise Stop will panic after a timeout informing that
// there is a session leak.
func (dbs *DBServer) Stop() {
	if dbs.session != nil {
		dbs.checkSessions()
		if dbs.session != nil {
			dbs.session.Close()
			dbs.session = nil
		}
	}
	if dbs.server != nil {
		dbs.tomb.Kill(nil)
		dbs.server.Process.Kill()
		select {
		case <-dbs.tomb.Dead():
		case <-time.After(5 * time.Second):
			panic("timeout waiting for mongod process to die")
		}
		dbs.server = nil
	}
}

// Session returns a new session to the server. The returned session
// must be closed after the test is done with it.
//
// The first Session obtained from a DBServer will start it.
func (dbs *DBServer) Session() *mgo.Session {
	if dbs.server == nil {
		dbs.start()
	}
	if dbs.session == nil {
		mgo.ResetStats()
		var err error
		dbs.session, err = mgo.Dial(dbs.host + "/test")
		if err != nil {
			panic(err)
		}
	}
	return dbs.session.Copy()
}

// checkSessions ensures all mgo sessions opened were properly closed.
// For slightly faster tests, it may be disabled setting the
// environmnet variable CHECK_SESSIONS to 0.
func (dbs *DBServer) checkSessions() {
	if check := os.Getenv("CHECK_SESSIONS"); check == "0" || dbs.server == nil || dbs.session == nil {
		return
	}
	dbs.session.Close()
	dbs.session = nil
	for i := 0; i < 100; i++ {
		stats := mgo.GetStats()
		if stats.SocketsInUse == 0 && stats.SocketsAlive == 0 {
			return
		}
		time.Sleep(100 * time.Millisecond)
	}
	panic("There are mgo sessions still alive.")
}

// Wipe drops all created databases and their data.
//
// The MongoDB server remains running if it was prevoiusly running,
// or stopped if it was previously stopped.
//
// All database sessions must be closed before or while the Wipe method
// is running. Otherwise Wipe will panic after a timeout informing that
// there is a session leak.
func (dbs *DBServer) Wipe() {
	if dbs.server == nil || dbs.session == nil {
		return
	}
	dbs.checkSessions()
	sessionUnset := dbs.session == nil
	session := dbs.Session()
	defer session.Close()
	if sessionUnset {
		dbs.session.Close()
		dbs.session = nil
	}
	names, err := session.DatabaseNames()
	if err != nil {
		panic(err)
	}
	for _, name := range names {
		switch name {
		case "admin", "local", "config":
		default:
			err = session.DB(name).DropDatabase()
			if err != nil {
				panic(err)
			}
		}
	}
}
<commit_msg>dbs.server.Process.Signal(os.Interrupt) instead of kill , kill causes mongod exit unexpected<commit_after>package dbtest

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"os/exec"
	"strconv"
	"time"

	"gopkg.in/mgo.v2-unstable"
	"gopkg.in/tomb.v2"
)

// DBServer controls a MongoDB server process to be used within test suites.
//
// The test server is started when Session is called the first time and should
// remain running for the duration of all tests, with the Wipe method being
// called between tests (before each of them) to clear stored data. After all tests
// are done, the Stop method should be called to stop the test server.
//
// Before the DBServer is used the SetPath method must be called to define
// the location for the database files to be stored.
type DBServer struct {
	session *mgo.Session
	output  bytes.Buffer
	server  *exec.Cmd
	dbpath  string
	host    string
	tomb    tomb.Tomb
}

// SetPath defines the path to the directory where the database files will be
// stored if it is started. The directory path itself is not created or removed
// by the test helper.
func (dbs *DBServer) SetPath(dbpath string) {
	dbs.dbpath = dbpath
}

func (dbs *DBServer) start() {
	if dbs.server != nil {
		panic("DBServer already started")
	}
	if dbs.dbpath == "" {
		panic("DBServer.SetPath must be called before using the server")
	}
	mgo.SetStats(true)
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		panic("unable to listen on a local address: " + err.Error())
	}
	addr := l.Addr().(*net.TCPAddr)
	l.Close()
	dbs.host = addr.String()

	args := []string{
		"--dbpath", dbs.dbpath,
		"--bind_ip", "127.0.0.1",
		"--port", strconv.Itoa(addr.Port),
		"--nssize", "1",
		"--noprealloc",
		"--smallfiles",
		"--nojournal",
	}
	dbs.tomb = tomb.Tomb{}
	dbs.server = exec.Command("mongod", args...)
	dbs.server.Stdout = &dbs.output
	dbs.server.Stderr = &dbs.output
	err = dbs.server.Start()
	if err != nil {
		panic(err)
	}
	dbs.tomb.Go(dbs.monitor)
	dbs.Wipe()
}

func (dbs *DBServer) monitor() error {
	dbs.server.Process.Wait()
	if dbs.tomb.Alive() {
		// Present some debugging information.
		fmt.Fprintf(os.Stderr, "---- mongod process died unexpectedly:\n")
		fmt.Fprintf(os.Stderr, "%s", dbs.output.Bytes())
		fmt.Fprintf(os.Stderr, "---- mongod processes running right now:\n")
		cmd := exec.Command("/bin/sh", "-c", "ps auxw | grep mongod")
		cmd.Stdout = os.Stderr
		cmd.Stderr = os.Stderr
		cmd.Run()
		fmt.Fprintf(os.Stderr, "----------------------------------------\n")

		panic("mongod process died unexpectedly")
	}
	return nil
}

// Stop stops the test server process, if it is running.
//
// It's okay to call Stop multiple times. After the test server is
// stopped it cannot be restarted.
//
// All database sessions must be closed before or while the Stop method
// is running. Otherwise Stop will panic after a timeout informing that
// there is a session leak.
func (dbs *DBServer) Stop() {
	if dbs.session != nil {
		dbs.checkSessions()
		if dbs.session != nil {
			dbs.session.Close()
			dbs.session = nil
		}
	}
	if dbs.server != nil {
		dbs.tomb.Kill(nil)
		dbs.server.Process.Signal(os.Interrupt)
		select {
		case <-dbs.tomb.Dead():
		case <-time.After(5 * time.Second):
			panic("timeout waiting for mongod process to die")
		}
		dbs.server = nil
	}
}

// Session returns a new session to the server. The returned session
// must be closed after the test is done with it.
//
// The first Session obtained from a DBServer will start it.
func (dbs *DBServer) Session() *mgo.Session {
	if dbs.server == nil {
		dbs.start()
	}
	if dbs.session == nil {
		mgo.ResetStats()
		var err error
		dbs.session, err = mgo.Dial(dbs.host + "/test")
		if err != nil {
			panic(err)
		}
	}
	return dbs.session.Copy()
}

// checkSessions ensures all mgo sessions opened were properly closed.
// For slightly faster tests, it may be disabled setting the
// environmnet variable CHECK_SESSIONS to 0.
func (dbs *DBServer) checkSessions() {
	if check := os.Getenv("CHECK_SESSIONS"); check == "0" || dbs.server == nil || dbs.session == nil {
		return
	}
	dbs.session.Close()
	dbs.session = nil
	for i := 0; i < 100; i++ {
		stats := mgo.GetStats()
		if stats.SocketsInUse == 0 && stats.SocketsAlive == 0 {
			return
		}
		time.Sleep(100 * time.Millisecond)
	}
	panic("There are mgo sessions still alive.")
}

// Wipe drops all created databases and their data.
//
// The MongoDB server remains running if it was prevoiusly running,
// or stopped if it was previously stopped.
//
// All database sessions must be closed before or while the Wipe method
// is running. Otherwise Wipe will panic after a timeout informing that
// there is a session leak.
func (dbs *DBServer) Wipe() {
	if dbs.server == nil || dbs.session == nil {
		return
	}
	dbs.checkSessions()
	sessionUnset := dbs.session == nil
	session := dbs.Session()
	defer session.Close()
	if sessionUnset {
		dbs.session.Close()
		dbs.session = nil
	}
	names, err := session.DatabaseNames()
	if err != nil {
		panic(err)
	}
	for _, name := range names {
		switch name {
		case "admin", "local", "config":
		default:
			err = session.DB(name).DropDatabase()
			if err != nil {
				panic(err)
			}
		}
	}
}
<|endoftext|>


<commit_before>package grim

// Copyright 2015 MediaMath <http://www.mediamath.com>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

import (
	"bufio"
	"fmt"
	"io"
	"os/exec"
	"sync"
	"syscall"
	"time"
)

type eitherStringOrError struct {
	str string
	err error
}

func execute(env []string, workingDir string, execPath string, timeout time.Duration, args ...string) (*executeResult, error) {
	outputChan := make(chan string)

	res, err := executeWithOutputChan(outputChan, env, workingDir, execPath, timeout, args...)
	if err != nil {
		return nil, err
	}

	out := ""
	for line := range outputChan {
		out += fmt.Sprintf("%v\n", line)
	}

	res.Output = out

	return res, nil
}

func executeWithOutputChan(outputChan chan string, env []string, workingDir string, execPath string, timeout time.Duration, args ...string) (*executeResult, error) {

	startTime := time.Now()

	cmd := exec.Command(execPath, args...)
	cmd.Dir = workingDir
	cmd.Env = env

	var startErr error

	var wg sync.WaitGroup

	outReader, orErr := cmd.StdoutPipe()
	if orErr != nil {
		return nil, fmt.Errorf("error capturing stdout: %v", orErr)
	}

	errReader, erErr := cmd.StderrPipe()
	if erErr != nil {
		return nil, fmt.Errorf("error capturing stderr: %v", erErr)
	}

	wg.Add(2)
	go sendLines(outReader, outputChan, &wg)
	go sendLines(errReader, outputChan, &wg)
	go closeAfterDone(outputChan, &wg)

	startErr = cmd.Start()
	if startErr != nil {
		return nil, fmt.Errorf("error starting process: %v", startErr)
	}

	exitCode, err := killProcessOnTimeout(cmd, timeout)
	if err != nil {
		return nil, err
	}

	return &executeResult{
		StartTime:  startTime,
		EndTime:    time.Now(),
		SysTime:    cmd.ProcessState.SystemTime(),
		UserTime:   cmd.ProcessState.UserTime(),
		InitialEnv: cmd.Env,
		ExitCode:   exitCode,
	}, nil
}

// kills a cmd process based on config timeout settings
func killProcessOnTimeout(cmd *exec.Cmd, timeout time.Duration) (int, error) {
	var exitCode int
	// 1 deep channel for done
	done := make(chan error, 1)

	go func() {
		done <- cmd.Wait()
	}()

	select {
	case <-time.After(timeout):
		if err := cmd.Process.Kill(); err != nil {
			return 0, fmt.Errorf("Failed to kill process: %v", err)
		}
		<-done
	case err := <-done:
		if err != nil {
			exitCode, err = getExitCode(err)
			if err != nil {
				return 0, fmt.Errorf("Build Error: %v", err)
			}
		}
	}
	return exitCode, nil
}

// gets the exit code from error
func getExitCode(err error) (int, error) {
	var exitCode int
	if exitErr, ok := err.(*exec.ExitError); ok {
		if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
			exitCode = status.ExitStatus()
		} else {
			return 0, fmt.Errorf("Wrong Wait Status: %v", err)
		}
	} else {
		return 0, fmt.Errorf("Can not cast to ExitError: %v", err)
	}
	return exitCode, nil
}

func sendLines(rc io.ReadCloser, linesChan chan string, wg *sync.WaitGroup) {
	scanner := bufio.NewScanner(rc)
	for scanner.Scan() {
		linesChan <- scanner.Text()
	}
	wg.Done()
}

func closeAfterDone(outputChan chan string, wg *sync.WaitGroup) {
	wg.Wait()
	close(outputChan)
}
<commit_msg>dont use zero exit code on timeout<commit_after>package grim

// Copyright 2015 MediaMath <http://www.mediamath.com>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

import (
	"bufio"
	"fmt"
	"io"
	"os/exec"
	"sync"
	"syscall"
	"time"
)

type eitherStringOrError struct {
	str string
	err error
}

func execute(env []string, workingDir string, execPath string, timeout time.Duration, args ...string) (*executeResult, error) {
	outputChan := make(chan string)

	res, err := executeWithOutputChan(outputChan, env, workingDir, execPath, timeout, args...)
	if err != nil {
		return nil, err
	}

	out := ""
	for line := range outputChan {
		out += fmt.Sprintf("%v\n", line)
	}

	res.Output = out

	return res, nil
}

func executeWithOutputChan(outputChan chan string, env []string, workingDir string, execPath string, timeout time.Duration, args ...string) (*executeResult, error) {

	startTime := time.Now()

	cmd := exec.Command(execPath, args...)
	cmd.Dir = workingDir
	cmd.Env = env

	var startErr error

	var wg sync.WaitGroup

	outReader, orErr := cmd.StdoutPipe()
	if orErr != nil {
		return nil, fmt.Errorf("error capturing stdout: %v", orErr)
	}

	errReader, erErr := cmd.StderrPipe()
	if erErr != nil {
		return nil, fmt.Errorf("error capturing stderr: %v", erErr)
	}

	wg.Add(2)
	go sendLines(outReader, outputChan, &wg)
	go sendLines(errReader, outputChan, &wg)
	go closeAfterDone(outputChan, &wg)

	startErr = cmd.Start()
	if startErr != nil {
		return nil, fmt.Errorf("error starting process: %v", startErr)
	}

	exitCode, err := killProcessOnTimeout(cmd, timeout)
	if err != nil {
		return nil, err
	}

	return &executeResult{
		StartTime:  startTime,
		EndTime:    time.Now(),
		SysTime:    cmd.ProcessState.SystemTime(),
		UserTime:   cmd.ProcessState.UserTime(),
		InitialEnv: cmd.Env,
		ExitCode:   exitCode,
	}, nil
}

// kills a cmd process based on config timeout settings
func killProcessOnTimeout(cmd *exec.Cmd, timeout time.Duration) (int, error) {
	var exitCode int
	// 1 deep channel for done
	done := make(chan error, 1)

	go func() {
		done <- cmd.Wait()
	}()

	select {
	case <-time.After(timeout):
		if err := cmd.Process.Kill(); err != nil {
			return 0, fmt.Errorf("Failed to kill process: %v", err)
		}
		<-done
		exitCode = -23
	case err := <-done:
		if err != nil {
			exitCode, err = getExitCode(err)
			if err != nil {
				return 0, fmt.Errorf("Build Error: %v", err)
			}
		}
	}
	return exitCode, nil
}

// gets the exit code from error
func getExitCode(err error) (int, error) {
	var exitCode int
	if exitErr, ok := err.(*exec.ExitError); ok {
		if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
			exitCode = status.ExitStatus()
		} else {
			return 0, fmt.Errorf("Wrong Wait Status: %v", err)
		}
	} else {
		return 0, fmt.Errorf("Can not cast to ExitError: %v", err)
	}
	return exitCode, nil
}

func sendLines(rc io.ReadCloser, linesChan chan string, wg *sync.WaitGroup) {
	scanner := bufio.NewScanner(rc)
	for scanner.Scan() {
		linesChan <- scanner.Text()
	}
	wg.Done()
}

func closeAfterDone(outputChan chan string, wg *sync.WaitGroup) {
	wg.Wait()
	close(outputChan)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"net/rpc"
	"testing"
)

type Args struct {
	C string
}

type Foo int

func (t *Foo) Dummy(args *Args, reply *string) error {
	*reply = "hello " + args.C
	return nil
}

var start = false

func BenchmarkHttpSync(b *testing.B) {
	done := make(chan bool, 10)
	fmt.Println("start")
	if !start {
		startHttpRPC()
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for i := 0; i < 10; i++ {
			go clientSync(done)
		}
		for i := 0; i < 10; i++ {
			<-done
		}
	}
}

func startHttpRPC() {
	start = true
	foo := new(Foo)

	rpc.Register(foo)
	rpc.HandleHTTP()
	l, e := net.Listen("tcp", ":1234")
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

func clientHttpDial() *rpc.Client {
	client, err := rpc.DialHTTP("tcp", "localhost:1234")
	if err != nil {
		log.Fatal("dialing:", err)
	}
	return client
}

func clientSync(done chan bool) {
	client := clientHttpDial()
	defer client.Close()

	args := &Args{"yifan"}
	var reply string

	err := client.Call("Foo.Dummy", args, &reply)
	if err != nil {
		log.Fatal("Dummy error:", err)
	}
	fmt.Println(reply)
	done <- true
}

func clientAsync(done chan bool) {
	client := clientHttpDial()
	defer client.Close()

	args := &Args{"yifan"}
	var reply string

	call := client.Go("Foo.Dummy", args, &reply, nil)
	<-call.Done
	if call.Error != nil {
		log.Fatal("Dummy error:", call.Error)
	}
	fmt.Println(reply)
	done <- true
}
<commit_msg>delete verbose<commit_after>package main

import (
	"fmt"
	"log"
	//"net"
	"net/http"
	"net/rpc"
	"testing"
)

type Args struct {
	C string
}

type Foo int

func (t *Foo) Dummy(args *Args, reply *string) error {
	*reply = "hello " + args.C
	return nil
}

var start = false

func BenchmarkHttpSync(b *testing.B) {
	done := make(chan bool, 10)
	if !start {
		fmt.Println("start Http RPC")
		startHttpRPC()
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for i := 0; i < 10; i++ {
			go clientSync(done)
		}
		for i := 0; i < 10; i++ {
			<-done
		}
	}
	b.StopTimer()
}

//func BenchmarkHttpAsync(b *testing.B) {
//	done := make(chan bool, 10)
//	if !start {
//		//fmt.Println("start Http RPC")
//		startHttpRPC()
//	}
//
//	b.ResetTimer()
//	for i := 0; i < b.N; i++ {
//		for i := 0; i < 10; i++ {
//			go clientAsync(done)
//		}
//		for i := 0; i < 10; i++ {
//			<-done
//		}
//	}
//}

func startHttpRPC() {
	start = true
	foo := new(Foo)

	rpc.Register(foo)
	rpc.HandleHTTP()
	go http.ListenAndServe(":1234", nil)
}

func clientHttpDial() *rpc.Client {
	client, err := rpc.DialHTTP("tcp", "localhost:1234")
	if err != nil {
		log.Fatal("dialing:", err)
	}
	return client
}

func clientSync(done chan bool) {
	client := clientHttpDial()
	defer client.Close()

	args := &Args{"yifan"}
	var reply string

	err := client.Call("Foo.Dummy", args, &reply)
	if err != nil {
		log.Fatal("Dummy error:", err)
	}
	//fmt.Println(reply)
	done <- true
}

func clientAsync(done chan bool) {
	client := clientHttpDial()
	defer client.Close()

	args := &Args{"yifan"}
	var reply string

	call := client.Go("Foo.Dummy", args, &reply, nil)
	<-call.Done
	if call.Error != nil {
		log.Fatal("Dummy error:", call.Error)
	}
	//fmt.Println(reply)
	done <- true
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"errors"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3"
	"github.com/ungerik/go-rss"
)

// CREATE TABLE feeds (name TEXT PRIMARY KEY, url TEXT NOT NULL, dayOfWeek INTEGER NOT NULL,
// seconds INTEGER NOT NULL, lastTitle TEXT NOT NULL);

type updatedTitleMessage struct {
	Name  string
	Title string
}

// Flag specifications.
var dbFilename = flag.String("db_file", "feeds.db", "filename of database to use")
var logFilename = flag.String("log_file", "", "filename to log to")
var target = flag.String("target", "", "target directory to download to")
var checkInterval = flag.Int(
	"check_interval", 3600, "seconds between checks during normal operation")
var rapidCheckInterval = flag.Int(
	"rapid_check_interval", 60, "seconds between checks when we suspect there will be a new item")
var rapidCheckDuration = flag.Int(
	"rapid_check_duration", 3600, "seconds that we suspect there will be a new item")
var downloadDelay = flag.Int(
	"download_delay", 30, "seconds to wait before downloading the file")
var requestDelay = flag.Int(
	"request_delay", 5, "seconds to wait between requests")

var requestDelayTicker <-chan time.Time

func downloadUrl(url string) error {
	// Figure out the filename to download to.
	lastSeparatorIndex := strings.LastIndex(url, "/")
	if lastSeparatorIndex == -1 {
		return errors.New("malformed url (no slash!?)")
	}
	filename := url[lastSeparatorIndex+1:]
	if len(filename) == 0 {
		return errors.New("malformed url (no filename)")
	}
	filepath := path.Join(*target, filename)

	// Actually download it.
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	file, err := os.Create(filepath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, resp.Body)
	if err != nil {
		return err
	}
	return nil
}

func lastRapidStartTime(fromTime time.Time, dayOfWeek int, seconds int) time.Time {
	dayDiff := dayOfWeek - int(fromTime.Weekday())
	if dayDiff > 0 {
		dayDiff -= 7
	}

	if dayDiff == 0 {
		if fromTime.Before(time.Date(fromTime.Year(), fromTime.Month(), fromTime.Day(), 0, 0, seconds, 0, time.Local)) {
			dayDiff -= 7
		}
	}

	return time.Date(
		fromTime.Year(), fromTime.Month(), fromTime.Day()+dayDiff, 0, 0, seconds, 0, time.Local)
}

func nextRapidStartTime(fromTime time.Time, dayOfWeek int, seconds int) time.Time {
	return lastRapidStartTime(fromTime.AddDate(0, 0, 7), dayOfWeek, seconds)
}

func isRapid(fromTime time.Time, dayOfWeek int, seconds int) bool {
	rapidStartTime := lastRapidStartTime(fromTime, dayOfWeek, seconds)
	return fromTime.Equal(rapidStartTime) || (fromTime.After(rapidStartTime) && fromTime.Before(rapidStartTime.Add(time.Duration(*rapidCheckDuration)*time.Second)))
}

func watchFeed(
	messages chan updatedTitleMessage, name string, feedUrl string, dayOfWeek int, seconds int,
	lastTitle string) {
	log.Printf("[%s] Starting watch.", name)

	checkTime := time.Now()

	// Main loop.
	for {
		// Fetch RSS.
		<-requestDelayTicker
		log.Printf("[%s] Checking for new items.", name)
		feed, err := rss.Read(feedUrl)
		if err != nil {
			log.Printf("[%s] Error fetching RSS: %s", name, err)
		} else {
			// Download any new files.
			for i := 0; i < len(feed.Item); i++ {
				if feed.Item[i].Title == lastTitle {
					break
				}

				log.Printf("[%s] Fetching %s.", name, feed.Item[i].Title)
				go func(title string, url string) {
					if *downloadDelay > 0 {
						time.Sleep(time.Duration(*downloadDelay) * time.Second)
					}
					<-requestDelayTicker
					err := downloadUrl(url)
					if err != nil {
						log.Printf("[%s] Error fetching %s: %s", name, url, err)
					} else {
						log.Printf("[%s] Fetched %s.", name, title)
					}
				}(feed.Item[i].Title, feed.Item[i].Link)
			}

			// Update last seen title.
			if len(feed.Item) > 0 {
				newTitle := feed.Item[0].Title
				if lastTitle != newTitle {
					lastTitle = newTitle
					messages <- updatedTitleMessage{name, lastTitle}
				}
			}
		}

		// Determine next wait time & wait.
		var nextCheckTime time.Time
		if isRapid(checkTime, dayOfWeek, seconds) {
			nextCheckTime = checkTime.Add(time.Duration(*rapidCheckInterval) * time.Second)
		} else {
			nextCheckTime = checkTime.Add(time.Duration(*checkInterval) * time.Second)
		}

		nextRapidTime := nextRapidStartTime(checkTime, dayOfWeek, seconds)
		if nextCheckTime.After(nextRapidTime) {
			nextCheckTime = nextRapidTime
		}
		checkTime = nextCheckTime
		time.Sleep(checkTime.Sub(time.Now()))
	}
}

func main() {
	// Check flags.
	flag.Parse()
	if *target == "" {
		log.Fatal("--target is required.")
	}

	// Set up logging.
	if *logFilename != "" {
		logWriter, err := os.OpenFile(*logFilename, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
		if err != nil {
			log.Fatal("--log_file could not be opened")
		}

		log.SetOutput(logWriter)
		defer logWriter.Close()
	}

	log.Print("Starting rss-downloader.")
	requestDelayTicker = time.Tick(time.Duration(*requestDelay) * time.Second)

	// Connect to database.
	db, err := sql.Open("sqlite3", *dbFilename)
	if err != nil {
		log.Fatalf("Error opening database connection: %s", err)
	}
	defer db.Close()

	// Start watching.
	messages := make(chan updatedTitleMessage)
	rows, err := db.Query("SELECT name, url, dayOfWeek, seconds, lastTitle FROM feeds")
	if err != nil {
		log.Fatalf("Error reading RSS feeds: %s", err)
	}
	for rows.Next() {
		var name string
		var url string
		var dayOfWeek int
		var seconds int
		var lastTitle string

		if err := rows.Scan(&name, &url, &dayOfWeek, &seconds, &lastTitle); err != nil {
			log.Fatalf("Error reading RSS feeds: %s", err)
		}

		go watchFeed(messages, name, url, dayOfWeek, seconds, lastTitle)
	}
	if err := rows.Err(); err != nil {
		log.Fatalf("Error reading RSS feeds: %s", err)
	}

	for {
		msg := <-messages
		_, err := db.Exec(
			"UPDATE feeds SET lastTitle = ? WHERE name = ?", msg.Title, msg.Name)
		if err != nil {
			log.Printf("[%s] Error updating last title: %s", err)
		}
	}
}
<commit_msg>When starting, do not check immediately.<commit_after>package main

import (
	"database/sql"
	"errors"
	"flag"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"strings"
	"math"
	"time"

	_ "github.com/mattn/go-sqlite3"
	"github.com/ungerik/go-rss"
)

// CREATE TABLE feeds (name TEXT PRIMARY KEY, url TEXT NOT NULL, dayOfWeek INTEGER NOT NULL,
// seconds INTEGER NOT NULL, lastTitle TEXT NOT NULL);

type updatedTitleMessage struct {
	Name  string
	Title string
}

// Flag specifications.
var dbFilename = flag.String("db_file", "feeds.db", "filename of database to use")
var logFilename = flag.String("log_file", "", "filename to log to")
var target = flag.String("target", "", "target directory to download to")
var checkInterval = flag.Int(
	"check_interval", 3600, "seconds between checks during normal operation")
var rapidCheckInterval = flag.Int(
	"rapid_check_interval", 60, "seconds between checks when we suspect there will be a new item")
var rapidCheckDuration = flag.Int(
	"rapid_check_duration", 3600, "seconds that we suspect there will be a new item")
var downloadDelay = flag.Int(
	"download_delay", 30, "seconds to wait before downloading the file")
var requestDelay = flag.Int(
	"request_delay", 5, "seconds to wait between requests")

var requestDelayTicker <-chan time.Time

func downloadUrl(url string) error {
	// Figure out the filename to download to.
	lastSeparatorIndex := strings.LastIndex(url, "/")
	if lastSeparatorIndex == -1 {
		return errors.New("malformed url (no slash!?)")
	}
	filename := url[lastSeparatorIndex+1:]
	if len(filename) == 0 {
		return errors.New("malformed url (no filename)")
	}
	filepath := path.Join(*target, filename)

	// Actually download it.
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	file, err := os.Create(filepath)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = io.Copy(file, resp.Body)
	if err != nil {
		return err
	}
	return nil
}

func lastRapidStartTime(fromTime time.Time, dayOfWeek int, seconds int) time.Time {
	dayDiff := dayOfWeek - int(fromTime.Weekday())
	if dayDiff > 0 {
		dayDiff -= 7
	}

	if dayDiff == 0 {
		if fromTime.Before(time.Date(fromTime.Year(), fromTime.Month(), fromTime.Day(), 0, 0, seconds, 0, time.Local)) {
			dayDiff -= 7
		}
	}

	return time.Date(
		fromTime.Year(), fromTime.Month(), fromTime.Day()+dayDiff, 0, 0, seconds, 0, time.Local)
}

func nextRapidStartTime(fromTime time.Time, dayOfWeek int, seconds int) time.Time {
	return lastRapidStartTime(fromTime.AddDate(0, 0, 7), dayOfWeek, seconds)
}

func isRapid(fromTime time.Time, dayOfWeek int, seconds int) bool {
	rapidStartTime := lastRapidStartTime(fromTime, dayOfWeek, seconds)
	return fromTime.Equal(rapidStartTime) || (fromTime.After(rapidStartTime) && fromTime.Before(rapidStartTime.Add(time.Duration(*rapidCheckDuration)*time.Second)))
}

func nextCheckTime(lastCheckTime time.Time, dayOfWeek int, seconds int) time.Time {
	var nextCheckTime time.Time

	if isRapid(lastCheckTime, dayOfWeek, seconds) {
		nextCheckTime = lastCheckTime.Add(time.Duration(*rapidCheckInterval) * time.Second)
	} else {
		nextCheckTime = lastCheckTime.Add(time.Duration(*checkInterval) * time.Second)
	}

	nextRapidTime := nextRapidStartTime(lastCheckTime, dayOfWeek, seconds)
	if nextCheckTime.After(nextRapidTime) {
		nextCheckTime = nextRapidTime
	}

	return nextCheckTime
}

func firstCheckTime(startTime time.Time, dayOfWeek int, seconds int) time.Time {
	// Grab info from last rapid start time.
	baseTime := lastRapidStartTime(startTime, dayOfWeek, seconds)
	var currentCheckInterval float64
	if isRapid(startTime, dayOfWeek, seconds) {
		currentCheckInterval = float64(*rapidCheckInterval)
	} else {
		baseTime = baseTime.Add(time.Duration(*rapidCheckDuration) * time.Second)
		currentCheckInterval = float64(*checkInterval)
	}

	// Calculate next check time.
	secondsOffsetFromBase := startTime.Sub(baseTime).Seconds()
	nextCheckOffsetFromBase := currentCheckInterval * math.Ceil(secondsOffsetFromBase / currentCheckInterval)
	nextCheckTime := baseTime.Add(time.Duration(nextCheckOffsetFromBase) * time.Second)

	// Fixup check time if it happens to be after the next rapid start time.
	nextRapidTime := nextRapidStartTime(startTime, dayOfWeek, seconds)
	if nextCheckTime.After(nextRapidTime) {
		nextCheckTime = nextRapidTime
	}

	return nextCheckTime
}

func watchFeed(
	messages chan updatedTitleMessage, name string, feedUrl string, dayOfWeek int, seconds int,
	lastTitle string) {
	log.Printf("[%s] Starting watch.", name)

	checkTime := firstCheckTime(time.Now(), dayOfWeek, seconds)

	// Main loop.
	for {
		// Wait until the next check time.
		time.Sleep(checkTime.Sub(time.Now()))
		checkTime = nextCheckTime(checkTime, dayOfWeek, seconds)

		// Fetch RSS.
		<-requestDelayTicker
		log.Printf("[%s] Checking for new items.", name)
		feed, err := rss.Read(feedUrl)
		if err != nil {
			log.Printf("[%s] Error fetching RSS: %s", name, err)
		} else {
			// Download any new files.
			for i := 0; i < len(feed.Item); i++ {
				if feed.Item[i].Title == lastTitle {
					break
				}

				log.Printf("[%s] Fetching %s.", name, feed.Item[i].Title)
				go func(title string, url string) {
					if *downloadDelay > 0 {
						time.Sleep(time.Duration(*downloadDelay) * time.Second)
					}
					<-requestDelayTicker
					err := downloadUrl(url)
					if err != nil {
						log.Printf("[%s] Error fetching %s: %s", name, url, err)
					} else {
						log.Printf("[%s] Fetched %s.", name, title)
					}
				}(feed.Item[i].Title, feed.Item[i].Link)
			}

			// Update last seen title.
			if len(feed.Item) > 0 {
				newTitle := feed.Item[0].Title
				if lastTitle != newTitle {
					lastTitle = newTitle
					messages <- updatedTitleMessage{name, lastTitle}
				}
			}
		}
	}
}

func main() {
	// Check flags.
	flag.Parse()
	if *target == "" {
		log.Fatal("--target is required.")
	}

	// Set up logging.
	if *logFilename != "" {
		logWriter, err := os.OpenFile(*logFilename, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
		if err != nil {
			log.Fatal("--log_file could not be opened")
		}

		log.SetOutput(logWriter)
		defer logWriter.Close()
	}

	log.Print("Starting rss-downloader.")
	requestDelayTicker = time.Tick(time.Duration(*requestDelay) * time.Second)

	// Connect to database.
	db, err := sql.Open("sqlite3", *dbFilename)
	if err != nil {
		log.Fatalf("Error opening database connection: %s", err)
	}
	defer db.Close()

	// Start watching.
	messages := make(chan updatedTitleMessage)
	rows, err := db.Query("SELECT name, url, dayOfWeek, seconds, lastTitle FROM feeds")
	if err != nil {
		log.Fatalf("Error reading RSS feeds: %s", err)
	}
	for rows.Next() {
		var name string
		var url string
		var dayOfWeek int
		var seconds int
		var lastTitle string

		if err := rows.Scan(&name, &url, &dayOfWeek, &seconds, &lastTitle); err != nil {
			log.Fatalf("Error reading RSS feeds: %s", err)
		}

		go watchFeed(messages, name, url, dayOfWeek, seconds, lastTitle)
	}
	if err := rows.Err(); err != nil {
		log.Fatalf("Error reading RSS feeds: %s", err)
	}

	for {
		msg := <-messages
		_, err := db.Exec(
			"UPDATE feeds SET lastTitle = ? WHERE name = ?", msg.Title, msg.Name)
		if err != nil {
			log.Printf("[%s] Error updating last title: %s", err)
		}
	}
}
<|endoftext|>


<commit_before>package transit

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/hashicorp/vault/logical"
)

const (
	shared    = false
	exclusive = true
)

type lockManager struct {
	// A lock for each named key
	locks map[string]*sync.RWMutex

	// A mutex for the map itself
	locksMutex sync.RWMutex

	// If caching is enabled, the map of name to in-memory policy cache
	cache map[string]*Policy

	// Used for global locking, and as the cache map mutex
	globalMutex sync.RWMutex
}

func newLockManager(cacheDisabled bool) *lockManager {
	lm := &lockManager{
		locks: map[string]*sync.RWMutex{},
	}
	if !cacheDisabled {
		lm.cache = map[string]*Policy{}
	}
	return lm
}

func (lm *lockManager) CacheActive() bool {
	return lm.cache != nil
}

func (lm *lockManager) LockAll(name string) {
	lm.globalMutex.Lock()
	lm.LockPolicy(name, exclusive)
}

func (lm *lockManager) UnlockAll(name string) {
	lm.UnlockPolicy(name, exclusive)
	lm.globalMutex.Unlock()
}

func (lm *lockManager) LockPolicy(name string, writeLock bool) {
	lm.locksMutex.RLock()
	lock := lm.locks[name]
	if lock != nil {
		// We want to give this up before locking the lock, but it's safe --
		// the only time we ever write to a value in this map is the first time
		// we access the value, so it won't be changing out from under us
		lm.locksMutex.RUnlock()
		if writeLock {
			lock.Lock()
		} else {
			lock.RLock()
		}
		return
	}

	lm.locksMutex.RUnlock()
	lm.locksMutex.Lock()

	// Don't defer the unlock call because if we get a valid lock below we want
	// to release the lock mutex right away to avoid the possibility of
	// deadlock by trying to grab the second lock

	// Check to make sure it hasn't been created since
	lock = lm.locks[name]
	if lock != nil {
		lm.locksMutex.Unlock()
		if writeLock {
			lock.Lock()
		} else {
			lock.RLock()
		}
		return
	}

	lock = &sync.RWMutex{}
	lm.locks[name] = lock
	lm.locksMutex.Unlock()
	if writeLock {
		lock.Lock()
	} else {
		lock.RLock()
	}
}

func (lm *lockManager) UnlockPolicy(name string, writeLock bool) {
	lm.locksMutex.RLock()
	lock := lm.locks[name]
	lm.locksMutex.RUnlock()

	if writeLock {
		lock.Unlock()
	} else {
		lock.RUnlock()
	}
}

func (lm *lockManager) GetPolicy(storage logical.Storage, name string) (*Policy, bool, error) {
	p, lt, _, err := lm.getPolicyCommon(storage, name, false, false)
	return p, lt, err
}

func (lm *lockManager) GetPolicyUpsert(storage logical.Storage, name string, derived bool) (*Policy, bool, bool, error) {
	return lm.getPolicyCommon(storage, name, true, derived)
}

// When the function returns, a lock will be held on the policy if err == nil.
// The type of lock will be indicated by the return value. It is the caller's
// responsibility to unlock.
func (lm *lockManager) getPolicyCommon(storage logical.Storage, name string, upsert, derived bool) (p *Policy, lockType bool, upserted bool, err error) {
	// If we are using a cache, lock it now to avoid having to do really
	// complicated lock juggling as we call various functions. We'll also defer
	// the store into the cache.
	lockType = shared
	lm.LockPolicy(name, shared)

	if lm.CacheActive() {
		lm.globalMutex.RLock()
		p = lm.cache[name]
		if p != nil {
			lm.globalMutex.RUnlock()
			return
		}
		lm.globalMutex.RUnlock()

		// When we return, since we didn't have the policy in the cache, if
		// there was no error, write the value in.
		defer func() {
			lm.globalMutex.Lock()
			defer lm.globalMutex.Unlock()
			// Make sure a policy didn't appear. If so, it will only be set if
			// there was no error, so now just clear the error and return that
			// policy.
			exp := lm.cache[name]
			if exp != nil {
				upserted = false
				err = nil
				p = exp
				return
			}

			if err == nil {
				lm.cache[name] = p
			}
		}()
	}

	p, err = lm.getStoredPolicy(storage, name)
	if err != nil {
		lm.UnlockPolicy(name, shared)
		return
	}

	if p == nil {
		if !upsert {
			lm.UnlockPolicy(name, shared)
			return
		}

		// Get an exlusive lock; on success, check again to ensure that no
		// policy exists. Note that if we are using a cache we will already be
		// serializing this entire code path and it's currently the only one
		// that generates policies, so we don't need to check the cache here;
		// simply checking the disk again is sufficient.
		lm.UnlockPolicy(name, shared)
		lockType = exclusive
		lm.LockPolicy(name, exclusive)

		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			return
		}
		if p != nil {
			return
		}

		upserted = true

		p = &Policy{
			Name:       name,
			CipherMode: "aes-gcm",
			Derived:    derived,
		}
		if derived {
			p.KDFMode = kdfMode
		}

		err = p.rotate(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			p = nil
		}

		// We don't need to worry about upgrading since it will be a new policy
		return
	}

	if p.needsUpgrade() {
		lm.UnlockPolicy(name, shared)
		lockType = exclusive
		lm.LockPolicy(name, exclusive)

		// Reload the policy with the write lock to ensure we still need the upgrade
		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			return
		}
		if p == nil {
			defer lm.UnlockPolicy(name, exclusive)
			err = fmt.Errorf("error reloading policy for upgrade")
			return
		}

		if !p.needsUpgrade() {
			// Already happened, return the newly loaded policy
			return
		}

		err = p.upgrade(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
		}
	}

	return
}

func (lm *lockManager) DeletePolicy(storage logical.Storage, name string) error {
	lm.LockAll(name)
	defer lm.UnlockAll(name)

	var p *Policy
	var err error

	if lm.CacheActive() {
		p = lm.cache[name]
		if p == nil {
			return fmt.Errorf("could not delete policy; not found")
		}
	} else {
		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			return err
		}
		if p == nil {
			return fmt.Errorf("could not delete policy; not found")
		}
	}

	if !p.DeletionAllowed {
		return fmt.Errorf("deletion is not allowed for this policy")
	}

	err = storage.Delete("policy/" + name)
	if err != nil {
		return fmt.Errorf("error deleting policy %s: %s", name, err)
	}

	err = storage.Delete("archive/" + name)
	if err != nil {
		return fmt.Errorf("error deleting archive %s: %s", name, err)
	}

	if lm.CacheActive() {
		delete(lm.cache, name)
	}

	return nil
}

// When this function returns it's the responsibility of the caller to call
// UnlockPolicy if err is nil and policy is not nil
func (lm *lockManager) RefreshPolicy(storage logical.Storage, name string) (p *Policy, err error) {
	lm.LockPolicy(name, exclusive)

	if lm.CacheActive() {
		p = lm.cache[name]
		if p != nil {
			return
		}
		err = fmt.Errorf("could not refresh policy; not found")
		defer lm.UnlockPolicy(name, exclusive)
		return
	}

	p, err = lm.getStoredPolicy(storage, name)
	if err != nil {
		defer lm.UnlockPolicy(name, exclusive)
		return
	}

	if p == nil {
		err = fmt.Errorf("could not refresh policy; not found")
		defer lm.UnlockPolicy(name, exclusive)
	}

	if p.needsUpgrade() {
		err = p.upgrade(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
		}
	}

	return
}

func (lm *lockManager) getStoredPolicy(storage logical.Storage, name string) (*Policy, error) {
	// Check if the policy already exists
	raw, err := storage.Get("policy/" + name)
	if err != nil {
		return nil, err
	}
	if raw == nil {
		return nil, nil
	}

	// Decode the policy
	policy := &Policy{
		Keys: KeyEntryMap{},
	}
	err = json.Unmarshal(raw.Value, policy)
	if err != nil {
		return nil, err
	}

	return policy, nil
}
<commit_msg>Change use-hint of lockAll and lockPolicy<commit_after>package transit

import (
	"encoding/json"
	"fmt"
	"sync"

	"github.com/hashicorp/vault/logical"
)

const (
	shared    = false
	exclusive = true
)

type lockManager struct {
	// A lock for each named key
	locks map[string]*sync.RWMutex

	// A mutex for the map itself
	locksMutex sync.RWMutex

	// If caching is enabled, the map of name to in-memory policy cache
	cache map[string]*Policy

	// Used for global locking, and as the cache map mutex
	globalMutex sync.RWMutex
}

func newLockManager(cacheDisabled bool) *lockManager {
	lm := &lockManager{
		locks: map[string]*sync.RWMutex{},
	}
	if !cacheDisabled {
		lm.cache = map[string]*Policy{}
	}
	return lm
}

func (lm *lockManager) CacheActive() bool {
	return lm.cache != nil
}

func (lm *lockManager) lockAll(name string) {
	lm.globalMutex.Lock()
	lm.lockPolicy(name, exclusive)
}

func (lm *lockManager) UnlockAll(name string) {
	lm.UnlockPolicy(name, exclusive)
	lm.globalMutex.Unlock()
}

func (lm *lockManager) lockPolicy(name string, writeLock bool) {
	lm.locksMutex.RLock()
	lock := lm.locks[name]
	if lock != nil {
		// We want to give this up before locking the lock, but it's safe --
		// the only time we ever write to a value in this map is the first time
		// we access the value, so it won't be changing out from under us
		lm.locksMutex.RUnlock()
		if writeLock {
			lock.Lock()
		} else {
			lock.RLock()
		}
		return
	}

	lm.locksMutex.RUnlock()
	lm.locksMutex.Lock()

	// Don't defer the unlock call because if we get a valid lock below we want
	// to release the lock mutex right away to avoid the possibility of
	// deadlock by trying to grab the second lock

	// Check to make sure it hasn't been created since
	lock = lm.locks[name]
	if lock != nil {
		lm.locksMutex.Unlock()
		if writeLock {
			lock.Lock()
		} else {
			lock.RLock()
		}
		return
	}

	lock = &sync.RWMutex{}
	lm.locks[name] = lock
	lm.locksMutex.Unlock()
	if writeLock {
		lock.Lock()
	} else {
		lock.RLock()
	}
}

func (lm *lockManager) UnlockPolicy(name string, writeLock bool) {
	lm.locksMutex.RLock()
	lock := lm.locks[name]
	lm.locksMutex.RUnlock()

	if writeLock {
		lock.Unlock()
	} else {
		lock.RUnlock()
	}
}

func (lm *lockManager) GetPolicy(storage logical.Storage, name string) (*Policy, bool, error) {
	p, lt, _, err := lm.getPolicyCommon(storage, name, false, false)
	return p, lt, err
}

func (lm *lockManager) GetPolicyUpsert(storage logical.Storage, name string, derived bool) (*Policy, bool, bool, error) {
	return lm.getPolicyCommon(storage, name, true, derived)
}

// When the function returns, a lock will be held on the policy if err == nil.
// The type of lock will be indicated by the return value. It is the caller's
// responsibility to unlock.
func (lm *lockManager) getPolicyCommon(storage logical.Storage, name string, upsert, derived bool) (p *Policy, lockType bool, upserted bool, err error) {
	// If we are using a cache, lock it now to avoid having to do really
	// complicated lock juggling as we call various functions. We'll also defer
	// the store into the cache.
	lockType = shared
	lm.lockPolicy(name, shared)

	if lm.CacheActive() {
		lm.globalMutex.RLock()
		p = lm.cache[name]
		if p != nil {
			lm.globalMutex.RUnlock()
			return
		}
		lm.globalMutex.RUnlock()

		// When we return, since we didn't have the policy in the cache, if
		// there was no error, write the value in.
		defer func() {
			lm.globalMutex.Lock()
			defer lm.globalMutex.Unlock()
			// Make sure a policy didn't appear. If so, it will only be set if
			// there was no error, so now just clear the error and return that
			// policy.
			exp := lm.cache[name]
			if exp != nil {
				upserted = false
				err = nil
				p = exp
				return
			}

			if err == nil {
				lm.cache[name] = p
			}
		}()
	}

	p, err = lm.getStoredPolicy(storage, name)
	if err != nil {
		lm.UnlockPolicy(name, shared)
		return
	}

	if p == nil {
		if !upsert {
			lm.UnlockPolicy(name, shared)
			return
		}

		// Get an exlusive lock; on success, check again to ensure that no
		// policy exists. Note that if we are using a cache we will already be
		// serializing this entire code path and it's currently the only one
		// that generates policies, so we don't need to check the cache here;
		// simply checking the disk again is sufficient.
		lm.UnlockPolicy(name, shared)
		lockType = exclusive
		lm.lockPolicy(name, exclusive)

		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			return
		}
		if p != nil {
			return
		}

		upserted = true

		p = &Policy{
			Name:       name,
			CipherMode: "aes-gcm",
			Derived:    derived,
		}
		if derived {
			p.KDFMode = kdfMode
		}

		err = p.rotate(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			p = nil
		}

		// We don't need to worry about upgrading since it will be a new policy
		return
	}

	if p.needsUpgrade() {
		lm.UnlockPolicy(name, shared)
		lockType = exclusive
		lm.lockPolicy(name, exclusive)

		// Reload the policy with the write lock to ensure we still need the upgrade
		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
			return
		}
		if p == nil {
			defer lm.UnlockPolicy(name, exclusive)
			err = fmt.Errorf("error reloading policy for upgrade")
			return
		}

		if !p.needsUpgrade() {
			// Already happened, return the newly loaded policy
			return
		}

		err = p.upgrade(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
		}
	}

	return
}

func (lm *lockManager) DeletePolicy(storage logical.Storage, name string) error {
	lm.lockAll(name)
	defer lm.UnlockAll(name)

	var p *Policy
	var err error

	if lm.CacheActive() {
		p = lm.cache[name]
		if p == nil {
			return fmt.Errorf("could not delete policy; not found")
		}
	} else {
		p, err = lm.getStoredPolicy(storage, name)
		if err != nil {
			return err
		}
		if p == nil {
			return fmt.Errorf("could not delete policy; not found")
		}
	}

	if !p.DeletionAllowed {
		return fmt.Errorf("deletion is not allowed for this policy")
	}

	err = storage.Delete("policy/" + name)
	if err != nil {
		return fmt.Errorf("error deleting policy %s: %s", name, err)
	}

	err = storage.Delete("archive/" + name)
	if err != nil {
		return fmt.Errorf("error deleting archive %s: %s", name, err)
	}

	if lm.CacheActive() {
		delete(lm.cache, name)
	}

	return nil
}

// When this function returns it's the responsibility of the caller to call
// UnlockPolicy if err is nil and policy is not nil
func (lm *lockManager) RefreshPolicy(storage logical.Storage, name string) (p *Policy, err error) {
	lm.lockPolicy(name, exclusive)

	if lm.CacheActive() {
		p = lm.cache[name]
		if p != nil {
			return
		}
		err = fmt.Errorf("could not refresh policy; not found")
		defer lm.UnlockPolicy(name, exclusive)
		return
	}

	p, err = lm.getStoredPolicy(storage, name)
	if err != nil {
		defer lm.UnlockPolicy(name, exclusive)
		return
	}

	if p == nil {
		err = fmt.Errorf("could not refresh policy; not found")
		defer lm.UnlockPolicy(name, exclusive)
	}

	if p.needsUpgrade() {
		err = p.upgrade(storage)
		if err != nil {
			defer lm.UnlockPolicy(name, exclusive)
		}
	}

	return
}

func (lm *lockManager) getStoredPolicy(storage logical.Storage, name string) (*Policy, error) {
	// Check if the policy already exists
	raw, err := storage.Get("policy/" + name)
	if err != nil {
		return nil, err
	}
	if raw == nil {
		return nil, nil
	}

	// Decode the policy
	policy := &Policy{
		Keys: KeyEntryMap{},
	}
	err = json.Unmarshal(raw.Value, policy)
	if err != nil {
		return nil, err
	}

	return policy, nil
}
<|endoftext|>


<commit_before>/*
 *
 * Copyright 2016, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

package grpc

import (
	"fmt"
	"sync"

	"golang.org/x/net/context"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/naming"
	"google.golang.org/grpc/transport"
)

// Address represents a server the client connects to.
// This is the EXPERIMENTAL API and may be changed or extended in the future.
type Address struct {
	// Addr is the server address on which a connection will be established.
	Addr string
	// Metadata is the information associated with Addr, which may be used
	// to make load balancing decision. This is from the metadata attached
	// in the address updates from name resolver.
	Metadata interface{}
}

// BalancerGetOption configures a Get call.
// This is the EXPERIMENTAL API and may be changed or extended in the future.
type BalancerGetOptions struct {
	// BlockingWait specifies whether Get should block when there is no
	// connected address.
	BlockingWait bool
}

// Balancer chooses network addresses for RPCs.
// This is the EXPERIMENTAL API and may be changed or extended in the future.
type Balancer interface {
	// Start does the initialization work to bootstrap a Balancer. For example,
	// this function may start the name resolution and watch the updates.
	Start(target string) error
	// Up informs the Balancer that gRPC has a connection to the server at
	// addr. It returns down which is called once the connection to addr gets
	// lost or closed.
	Up(addr Address) (down func(error))
	// Get gets the address of a server for the RPC corresponding to ctx.
	// i) If it returns a connected address, gRPC internals issues the RPC on the
	// connection to this address;
	// ii) If it returns an address on which the connection is under construction
	// (initiated by Notify(...)) but not connected, gRPC internals
	//  * fails RPC if the RPC is fail-fast and connection is in the TransientFailure
	//  or Shutdown state;
	//  or
	//  * issues RPC on the connection otherwise.
	// iii) If it returns an address on which the connection does not exist, gRPC
	// internals treats it as an error and will fail the corresponding RPC.
	//
	// Therefore, we recommend the following general rule when you write your own
	// custom Balancer.
	//
	// If opts.BlockingWait is true, it should return a connected address or
	// block if there is no connected address. It respects the timeout or
	// cancellation of ctx when blocking. If opts.BlockingWait is false (for fail-fast
	// RPCs), it should return an address it has notified via Notify(...) immediately
	// instead of blocking.
	//
	// The function returns put which is called once the rpc has completed or failed.
	// put can collect and report rpc stats to a remote load balancer.
	Get(ctx context.Context, opts BalancerGetOptions) (addr Address, put func(), err error)
	// Notify notifies gRPC internals the list of Address to be connected. The list
	// may be from a name resolver or remote load balancer. gRPC internals will
	// compare it with the exisiting connected addresses. If the address Balancer
	// notified is not in the exisitng connected addresses, gRPC starts to connect
	// the address. If an address in the exisiting connected addresses is not in
	// the notification list, the corresponding connection is shutdown gracefully.
	// Otherwise, there are no operations to take. Note that this function must
	// return the full list of the Addrresses which should be connected. It is NOT delta.
	Notify() <-chan []Address
	// Close shuts down the balancer.
	Close() error
}

// downErr implements net.Error. It is contructed by gRPC internals and passed to the down
// call of Balancer.
type downErr struct {
	timeout   bool
	temporary bool
	desc      string
}

func (e downErr) Error() string   { return e.desc }
func (e downErr) Timeout() bool   { return e.timeout }
func (e downErr) Temporary() bool { return e.temporary }

func downErrorf(timeout, temporary bool, format string, a ...interface{}) downErr {
	return downErr{
		timeout:   timeout,
		temporary: temporary,
		desc:      fmt.Sprintf(format, a...),
	}
}

// RoundRobin returns a Balancer that selects addresses round-robin. It starts to watch
// the name resolution updates.
func RoundRobin(r naming.Resolver) Balancer {
	return &roundRobin{r: r}
}

type roundRobin struct {
	r         naming.Resolver
	open      []Address // all the addresses the client should potentially connect
	mu        sync.Mutex
	addrCh    chan []Address // the channel to notify gRPC internals the list of addresses the client should connect to.
	connected []Address      // all the connected addresses
	next      int            // index of the next address to return for Get()
	waitCh    chan struct{}  // the channel to block when there is no connected address available
	done      bool           // The Balancer is closed.
}

func (rr *roundRobin) watchAddrUpdates(w naming.Watcher) error {
	updates, err := w.Next()
	if err != nil {
		return err
	}
	rr.mu.Lock()
	defer rr.mu.Unlock()
	for _, update := range updates {
		addr := Address{
			Addr:     update.Addr,
			Metadata: update.Metadata,
		}
		switch update.Op {
		case naming.Add:
			var exist bool
			for _, v := range rr.open {
				if addr == v {
					exist = true
					grpclog.Println("grpc: The name resolver wanted to add an existing address: ", addr)
					break
				}
			}
			if exist {
				continue
			}
			rr.open = append(rr.open, addr)
		case naming.Delete:
			for i, v := range rr.open {
				if v == addr {
					copy(rr.open[i:], rr.open[i+1:])
					rr.open = rr.open[:len(rr.open)-1]
					break
				}
			}
		default:
			grpclog.Println("Unknown update.Op ", update.Op)
		}
	}
	// Make a copy of rr.open and write it onto rr.addrCh so that gRPC internals gets notified.
	open := make([]Address, len(rr.open))
	copy(open, rr.open)
	if rr.done {
		return ErrClientConnClosing
	}
	rr.addrCh <- open
	return nil
}

func (rr *roundRobin) Start(target string) error {
	if rr.r == nil {
		// If there is no name resolver installed, it is not needed to
		// do name resolution. In this case, rr.addrCh stays nil.
		return nil
	}
	w, err := rr.r.Resolve(target)
	if err != nil {
		return err
	}
	rr.addrCh = make(chan []Address)
	go func() {
		for {
			if err := rr.watchAddrUpdates(w); err != nil {
				return
			}
		}
	}()
	return nil
}

// Up appends addr to the end of rr.addrs and sends notification if there
// are pending Get() calls.
func (rr *roundRobin) Up(addr Address) func(error) {
	rr.mu.Lock()
	defer rr.mu.Unlock()
	for _, a := range rr.connected {
		if a == addr {
			return nil
		}
	}
	rr.connected = append(rr.connected, addr)
	if len(rr.connected) == 1 {
		// addr is only one available. Notify the Get() callers who are blocking.
		if rr.waitCh != nil {
			close(rr.waitCh)
			rr.waitCh = nil
		}
	}
	return func(err error) {
		rr.down(addr, err)
	}
}

// down removes addr from rr.addrs and moves the remaining addrs forward.
func (rr *roundRobin) down(addr Address, err error) {
	rr.mu.Lock()
	defer rr.mu.Unlock()
	for i, a := range rr.connected {
		if a == addr {
			copy(rr.connected[i:], rr.connected[i+1:])
			rr.connected = rr.connected[:len(rr.connected)-1]
			return
		}
	}
}

// Get returns the next addr in the rotation.
func (rr *roundRobin) Get(ctx context.Context, opts BalancerGetOptions) (addr Address, put func(), err error) {
	var ch chan struct{}
	rr.mu.Lock()
	if rr.done {
		rr.mu.Unlock()
		err = ErrClientConnClosing
		return
	}
	if rr.next >= len(rr.connected) {
		rr.next = 0
	}
	if len(rr.connected) > 0 {
		addr = rr.connected[rr.next]
		rr.next++
		rr.mu.Unlock()
		return
	}
	// There is no address available. Wait on rr.waitCh.
	// TODO(zhaoq): Handle the case when opts.BlockingWait is false.
	if rr.waitCh == nil {
		ch = make(chan struct{})
		rr.waitCh = ch
	} else {
		ch = rr.waitCh
	}
	rr.mu.Unlock()
	for {
		select {
		case <-ctx.Done():
			err = transport.ContextErr(ctx.Err())
			return
		case <-ch:
			rr.mu.Lock()
			if rr.done {
				rr.mu.Unlock()
				err = ErrClientConnClosing
				return
			}
			if len(rr.connected) == 0 {
				// The newly added addr got removed by Down() again.
				rr.mu.Unlock()
				continue
			}
			if rr.next >= len(rr.connected) {
				rr.next = 0
			}
			addr = rr.connected[rr.next]
			rr.next++
			rr.mu.Unlock()
			return
		}
	}
}

func (rr *roundRobin) Notify() <-chan []Address {
	return rr.addrCh
}

func (rr *roundRobin) Close() error {
	rr.mu.Lock()
	defer rr.mu.Unlock()
	rr.done = true
	if rr.waitCh != nil {
		close(rr.waitCh)
		rr.waitCh = nil
	}
	if rr.addrCh != nil {
		close(rr.addrCh)
	}
	return nil
}
<commit_msg>fix some typos<commit_after>/*
 *
 * Copyright 2016, Google Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *
 *     * Redistributions of source code must retain the above copyright
 * notice, this list of conditions and the following disclaimer.
 *     * Redistributions in binary form must reproduce the above
 * copyright notice, this list of conditions and the following disclaimer
 * in the documentation and/or other materials provided with the
 * distribution.
 *     * Neither the name of Google Inc. nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

package grpc

import (
	"fmt"
	"sync"

	"golang.org/x/net/context"
	"google.golang.org/grpc/grpclog"
	"google.golang.org/grpc/naming"
	"google.golang.org/grpc/transport"
)

// Address represents a server the client connects to.
// This is the EXPERIMENTAL API and may be changed or extended in the future.
type Address struct {
	// Addr is the server address on which a connection will be established.
	Addr string
	// Metadata is the information associated with Addr, which may be used
	// to make load balancing decision. This is from the metadata attached
	// in the address updates from name resolver.
	Metadata interface{}
}

// BalancerGetOptions configures a Get call.
// This is the EXPERIMENTAL API and may be changed or extended in the future.
type BalancerGetOptions struct {
	// BlockingWait specifies whether Get should block when there is no
	// connected address.
	BlockingWait bool
}

// Balancer chooses network addresses for RPCs.
// This is the EXPERIMENTAL API and may be changed or extended in the future.
type Balancer interface {
	// Start does the initialization work to bootstrap a Balancer. For example,
	// this function may start the name resolution and watch the updates.
	Start(target string) error
	// Up informs the Balancer that gRPC has a connection to the server at
	// addr. It returns down which is called once the connection to addr gets
	// lost or closed.
	Up(addr Address) (down func(error))
	// Get gets the address of a server for the RPC corresponding to ctx.
	// i) If it returns a connected address, gRPC internals issues the RPC on the
	// connection to this address;
	// ii) If it returns an address on which the connection is under construction
	// (initiated by Notify(...)) but not connected, gRPC internals
	//  * fails RPC if the RPC is fail-fast and connection is in the TransientFailure
	//  or Shutdown state;
	//  or
	//  * issues RPC on the connection otherwise.
	// iii) If it returns an address on which the connection does not exist, gRPC
	// internals treats it as an error and will fail the corresponding RPC.
	//
	// Therefore, we recommend the following general rule when you write your own
	// custom Balancer.
	//
	// If opts.BlockingWait is true, it should return a connected address or
	// block if there is no connected address. It respects the timeout or
	// cancellation of ctx when blocking. If opts.BlockingWait is false (for fail-fast
	// RPCs), it should return an address it has notified via Notify(...) immediately
	// instead of blocking.
	//
	// The function returns put which is called once the rpc has completed or failed.
	// put can collect and report RPC stats to a remote load balancer.
	Get(ctx context.Context, opts BalancerGetOptions) (addr Address, put func(), err error)
	// Notify notifies gRPC internals the list of Address to be connected. The list
	// may be from a name resolver or remote load balancer. gRPC internals will
	// compare it with the existing connected addresses. If the address Balancer
	// notified is not in the existing connected addresses, gRPC starts to connect
	// the address. If an address in the existing connected addresses is not in
	// the notification list, the corresponding connection is shutdown gracefully.
	// Otherwise, there are no operations to take. Note that this function must
	// return the full list of the Addresses which should be connected. It is NOT delta.
	Notify() <-chan []Address
	// Close shuts down the balancer.
	Close() error
}

// downErr implements net.Error. It is constructed by gRPC internals and passed to the down
// call of Balancer.
type downErr struct {
	timeout   bool
	temporary bool
	desc      string
}

func (e downErr) Error() string   { return e.desc }
func (e downErr) Timeout() bool   { return e.timeout }
func (e downErr) Temporary() bool { return e.temporary }

func downErrorf(timeout, temporary bool, format string, a ...interface{}) downErr {
	return downErr{
		timeout:   timeout,
		temporary: temporary,
		desc:      fmt.Sprintf(format, a...),
	}
}

// RoundRobin returns a Balancer that selects addresses round-robin. It starts to watch
// the name resolution updates.
func RoundRobin(r naming.Resolver) Balancer {
	return &roundRobin{r: r}
}

type roundRobin struct {
	r         naming.Resolver
	open      []Address // all the addresses the client should potentially connect
	mu        sync.Mutex
	addrCh    chan []Address // the channel to notify gRPC internals the list of addresses the client should connect to.
	connected []Address      // all the connected addresses
	next      int            // index of the next address to return for Get()
	waitCh    chan struct{}  // the channel to block when there is no connected address available
	done      bool           // The Balancer is closed.
}

func (rr *roundRobin) watchAddrUpdates(w naming.Watcher) error {
	updates, err := w.Next()
	if err != nil {
		return err
	}
	rr.mu.Lock()
	defer rr.mu.Unlock()
	for _, update := range updates {
		addr := Address{
			Addr:     update.Addr,
			Metadata: update.Metadata,
		}
		switch update.Op {
		case naming.Add:
			var exist bool
			for _, v := range rr.open {
				if addr == v {
					exist = true
					grpclog.Println("grpc: The name resolver wanted to add an existing address: ", addr)
					break
				}
			}
			if exist {
				continue
			}
			rr.open = append(rr.open, addr)
		case naming.Delete:
			for i, v := range rr.open {
				if v == addr {
					copy(rr.open[i:], rr.open[i+1:])
					rr.open = rr.open[:len(rr.open)-1]
					break
				}
			}
		default:
			grpclog.Println("Unknown update.Op ", update.Op)
		}
	}
	// Make a copy of rr.open and write it onto rr.addrCh so that gRPC internals gets notified.
	open := make([]Address, len(rr.open))
	copy(open, rr.open)
	if rr.done {
		return ErrClientConnClosing
	}
	rr.addrCh <- open
	return nil
}

func (rr *roundRobin) Start(target string) error {
	if rr.r == nil {
		// If there is no name resolver installed, it is not needed to
		// do name resolution. In this case, rr.addrCh stays nil.
		return nil
	}
	w, err := rr.r.Resolve(target)
	if err != nil {
		return err
	}
	rr.addrCh = make(chan []Address)
	go func() {
		for {
			if err := rr.watchAddrUpdates(w); err != nil {
				return
			}
		}
	}()
	return nil
}

// Up appends addr to the end of rr.addrs and sends notification if there
// are pending Get() calls.
func (rr *roundRobin) Up(addr Address) func(error) {
	rr.mu.Lock()
	defer rr.mu.Unlock()
	for _, a := range rr.connected {
		if a == addr {
			return nil
		}
	}
	rr.connected = append(rr.connected, addr)
	if len(rr.connected) == 1 {
		// addr is only one available. Notify the Get() callers who are blocking.
		if rr.waitCh != nil {
			close(rr.waitCh)
			rr.waitCh = nil
		}
	}
	return func(err error) {
		rr.down(addr, err)
	}
}

// down removes addr from rr.addrs and moves the remaining addrs forward.
func (rr *roundRobin) down(addr Address, err error) {
	rr.mu.Lock()
	defer rr.mu.Unlock()
	for i, a := range rr.connected {
		if a == addr {
			copy(rr.connected[i:], rr.connected[i+1:])
			rr.connected = rr.connected[:len(rr.connected)-1]
			return
		}
	}
}

// Get returns the next addr in the rotation.
func (rr *roundRobin) Get(ctx context.Context, opts BalancerGetOptions) (addr Address, put func(), err error) {
	var ch chan struct{}
	rr.mu.Lock()
	if rr.done {
		rr.mu.Unlock()
		err = ErrClientConnClosing
		return
	}
	if rr.next >= len(rr.connected) {
		rr.next = 0
	}
	if len(rr.connected) > 0 {
		addr = rr.connected[rr.next]
		rr.next++
		rr.mu.Unlock()
		return
	}
	// There is no address available. Wait on rr.waitCh.
	// TODO(zhaoq): Handle the case when opts.BlockingWait is false.
	if rr.waitCh == nil {
		ch = make(chan struct{})
		rr.waitCh = ch
	} else {
		ch = rr.waitCh
	}
	rr.mu.Unlock()
	for {
		select {
		case <-ctx.Done():
			err = transport.ContextErr(ctx.Err())
			return
		case <-ch:
			rr.mu.Lock()
			if rr.done {
				rr.mu.Unlock()
				err = ErrClientConnClosing
				return
			}
			if len(rr.connected) == 0 {
				// The newly added addr got removed by Down() again.
				rr.mu.Unlock()
				continue
			}
			if rr.next >= len(rr.connected) {
				rr.next = 0
			}
			addr = rr.connected[rr.next]
			rr.next++
			rr.mu.Unlock()
			return
		}
	}
}

func (rr *roundRobin) Notify() <-chan []Address {
	return rr.addrCh
}

func (rr *roundRobin) Close() error {
	rr.mu.Lock()
	defer rr.mu.Unlock()
	rr.done = true
	if rr.waitCh != nil {
		close(rr.waitCh)
		rr.waitCh = nil
	}
	if rr.addrCh != nil {
		close(rr.addrCh)
	}
	return nil
}
<|endoftext|>


<commit_before>package diff

import (
	"bytes"
	"fmt"
	"io"
	"math"
	"sort"
	"strings"

	"github.com/aryann/difflib"
	"github.com/mgutz/ansi"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime/serializer/json"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/client-go/kubernetes/scheme"

	"github.com/databus23/helm-diff/v3/manifest"
)

// Manifests diff on manifests
func Manifests(oldIndex, newIndex map[string]*manifest.MappingResult, suppressedKinds []string, showSecrets bool, context int, output string, stripTrailingCR bool, to io.Writer) bool {
	report := Report{}
	report.setupReportFormat(output)
	seenAnyChanges := false
	emptyMapping := &manifest.MappingResult{}
	for _, key := range sortedKeys(oldIndex) {
		oldContent := oldIndex[key]

		if newContent, ok := newIndex[key]; ok {
			if oldContent.Content != newContent.Content {
				// modified
				if !showSecrets {
					redactSecrets(oldContent, newContent)
				}

				diffs := diffMappingResults(oldContent, newContent, stripTrailingCR)
				if len(diffs) > 0 {
					seenAnyChanges = true
				}
				report.addEntry(key, suppressedKinds, oldContent.Kind, context, diffs, "MODIFY")
			}
		} else {
			// removed
			if !showSecrets {
				redactSecrets(oldContent, nil)

			}
			diffs := diffMappingResults(oldContent, emptyMapping, stripTrailingCR)
			if len(diffs) > 0 {
				seenAnyChanges = true
			}
			report.addEntry(key, suppressedKinds, oldContent.Kind, context, diffs, "REMOVE")
		}
	}

	for _, key := range sortedKeys(newIndex) {
		newContent := newIndex[key]

		if _, ok := oldIndex[key]; !ok {
			// added
			if !showSecrets {
				redactSecrets(nil, newContent)
			}
			diffs := diffMappingResults(emptyMapping, newContent, stripTrailingCR)
			if len(diffs) > 0 {
				seenAnyChanges = true
			}
			report.addEntry(key, suppressedKinds, newContent.Kind, context, diffs, "ADD")
		}
	}
	report.print(to)
	report.clean()
	return seenAnyChanges
}

func redactSecrets(old, new *manifest.MappingResult) {
	if (old != nil && old.Kind != "Secret") || (new != nil && new.Kind != "Secret") {
		return
	}
	serializer := json.NewYAMLSerializer(json.DefaultMetaFactory, scheme.Scheme,
		scheme.Scheme)
	var oldSecret, newSecret v1.Secret

	if old != nil {
		if err := yaml.NewYAMLToJSONDecoder(bytes.NewBufferString(old.Content)).Decode(&oldSecret); err != nil {
			old.Content = fmt.Sprintf("Error parsing old secret: %s", err)
		}
	}
	if new != nil {
		if err := yaml.NewYAMLToJSONDecoder(bytes.NewBufferString(new.Content)).Decode(&newSecret); err != nil {
			new.Content = fmt.Sprintf("Error parsing new secret: %s", err)
		}
	}
	if old != nil {
		oldSecret.StringData = make(map[string]string, len(oldSecret.Data))
		for k, v := range oldSecret.Data {
			if new != nil && bytes.Equal(v, newSecret.Data[k]) {
				oldSecret.StringData[k] = fmt.Sprintf("REDACTED # (%d bytes)", len(v))
			} else {
				oldSecret.StringData[k] = fmt.Sprintf("-------- # (%d bytes)", len(v))
			}
		}
	}
	if new != nil {
		newSecret.StringData = make(map[string]string, len(newSecret.Data))
		for k, v := range newSecret.Data {
			if old != nil && bytes.Equal(v, oldSecret.Data[k]) {
				newSecret.StringData[k] = fmt.Sprintf("REDACTED # (%d bytes)", len(v))
			} else {
				newSecret.StringData[k] = fmt.Sprintf("++++++++ # (%d bytes)", len(v))
			}
		}
	}
	// remove Data field now that we are using StringData for serialization
	var buf bytes.Buffer
	if old != nil {
		oldSecret.Data = nil
		if err := serializer.Encode(&oldSecret, &buf); err != nil {
			new.Content = fmt.Sprintf("Error encoding new secret: %s", err)
		}
		old.Content = getComment(old.Content) + strings.Replace(strings.Replace(buf.String(), "stringData", "data", 1), "  creationTimestamp: null\n", "", 1)
		buf.Reset() //reuse buffer for new secret
	}
	if new != nil {
		newSecret.Data = nil
		if err := serializer.Encode(&newSecret, &buf); err != nil {
			new.Content = fmt.Sprintf("Error encoding new secret: %s", err)
		}
		new.Content = getComment(new.Content) + strings.Replace(strings.Replace(buf.String(), "stringData", "data", 1), "  creationTimestamp: null\n", "", 1)
	}
}

// return the first line of a string if its a comment.
// This gives as the # Source: lines from the rendering
func getComment(s string) string {
	i := strings.Index(s, "\n")
	if i < 0 || !strings.HasPrefix(s, "#") {
		return ""
	}
	return s[:i+1]

}

// Releases reindex the content  based on the template names and pass it to Manifests
func Releases(oldIndex, newIndex map[string]*manifest.MappingResult, suppressedKinds []string, showSecrets bool, context int, output string, stripTrailingCR bool, to io.Writer) bool {
	oldIndex = reIndexForRelease(oldIndex)
	newIndex = reIndexForRelease(newIndex)
	return Manifests(oldIndex, newIndex, suppressedKinds, showSecrets, context, output, stripTrailingCR, to)
}

func diffMappingResults(oldContent *manifest.MappingResult, newContent *manifest.MappingResult, stripTrailingCR bool) []difflib.DiffRecord {
	return diffStrings(oldContent.Content, newContent.Content, stripTrailingCR)
}

func diffStrings(before, after string, stripTrailingCR bool) []difflib.DiffRecord {
	return difflib.Diff(split(before, stripTrailingCR), split(after, stripTrailingCR))
}

func split(value string, stripTrailingCR bool) []string {
	const sep = "\n"
	split := strings.Split(value, sep)
	if !stripTrailingCR {
		return split
	}
	var stripped []string
	for _, s := range split {
		stripped = append(stripped, strings.TrimSuffix(s, "\r"))
	}
	return stripped
}

func printDiffRecords(suppressedKinds []string, kind string, context int, diffs []difflib.DiffRecord, to io.Writer) {
	for _, ckind := range suppressedKinds {

		if ckind == kind {
			str := fmt.Sprintf("+ Changes suppressed on sensitive content of type %s\n", kind)
			fmt.Fprintf(to, ansi.Color(str, "yellow"))
			return
		}
	}

	if context >= 0 {
		distances := calculateDistances(diffs)
		omitting := false
		for i, diff := range diffs {
			if distances[i] > context {
				if !omitting {
					fmt.Fprintln(to, "...")
					omitting = true
				}
			} else {
				omitting = false
				printDiffRecord(diff, to)
			}
		}
	} else {
		for _, diff := range diffs {
			printDiffRecord(diff, to)
		}
	}
}

func printDiffRecord(diff difflib.DiffRecord, to io.Writer) {
	text := diff.Payload

	switch diff.Delta {
	case difflib.RightOnly:
		fmt.Fprintf(to, "%s\n", ansi.Color("+ "+text, "green"))
	case difflib.LeftOnly:
		fmt.Fprintf(to, "%s\n", ansi.Color("- "+text, "red"))
	case difflib.Common:
		if text == "" {
			fmt.Fprintln(to)
		} else {
			fmt.Fprintf(to, "%s\n", "  "+text)
		}
	}
}

// Calculate distance of every diff-line to the closest change
func calculateDistances(diffs []difflib.DiffRecord) map[int]int {
	distances := map[int]int{}

	// Iterate forwards through diffs, set 'distance' based on closest 'change' before this line
	change := -1
	for i, diff := range diffs {
		if diff.Delta != difflib.Common {
			change = i
		}
		distance := math.MaxInt32
		if change != -1 {
			distance = i - change
		}
		distances[i] = distance
	}

	// Iterate backwards through diffs, reduce 'distance' based on closest 'change' after this line
	change = -1
	for i := len(diffs) - 1; i >= 0; i-- {
		diff := diffs[i]
		if diff.Delta != difflib.Common {
			change = i
		}
		if change != -1 {
			distance := change - i
			if distance < distances[i] {
				distances[i] = distance
			}
		}
	}

	return distances
}

// reIndexForRelease based on template names
func reIndexForRelease(index map[string]*manifest.MappingResult) map[string]*manifest.MappingResult {

	// sort the index to iterate map in the same order
	var keys []string
	for key := range index {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	// holds number of object in a single file
	count := make(map[string]int)

	newIndex := make(map[string]*manifest.MappingResult)

	for key := range keys {

		str := strings.Replace(strings.Split(index[keys[key]].Content, "\n")[0], "# Source: ", "", 1)

		if _, ok := newIndex[str]; ok {
			count[str]++
			str += fmt.Sprintf(" %d", count[str])
			newIndex[str] = index[keys[key]]
		} else {
			newIndex[str] = index[keys[key]]
			count[str]++
		}
	}
	return newIndex
}

func sortedKeys(manifests map[string]*manifest.MappingResult) []string {
	var keys []string

	for key := range manifests {
		keys = append(keys, key)
	}

	sort.Strings(keys)

	return keys
}
<commit_msg>Extract redaction/diff/report generation<commit_after>package diff

import (
	"bytes"
	"fmt"
	"io"
	"math"
	"sort"
	"strings"

	"github.com/aryann/difflib"
	"github.com/mgutz/ansi"
	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime/serializer/json"
	"k8s.io/apimachinery/pkg/util/yaml"
	"k8s.io/client-go/kubernetes/scheme"

	"github.com/databus23/helm-diff/v3/manifest"
)

// Manifests diff on manifests
func Manifests(oldIndex, newIndex map[string]*manifest.MappingResult, suppressedKinds []string, showSecrets bool, context int, output string, stripTrailingCR bool, to io.Writer) bool {
	report := Report{}
	report.setupReportFormat(output)
	for _, key := range sortedKeys(oldIndex) {
		oldContent := oldIndex[key]

		if newContent, ok := newIndex[key]; ok {
			// modified?
			doDiff(key, oldContent, newContent, suppressedKinds, showSecrets, stripTrailingCR, context)
		} else {
			// removed
			doDiff(key, oldContent, nil, suppressedKinds, showSecrets, stripTrailingCR, context)
		}
	}

	for _, key := range sortedKeys(newIndex) {
		newContent := newIndex[key]

		if _, ok := oldIndex[key]; !ok {
			doDiff(key, nil, newContent, suppressedKinds, showSecrets, stripTrailingCR, context)
		}
	}

	seenAnyChanges := len(report.entries) > 0
	report.print(to)
	report.clean()
	return seenAnyChanges
}

func actualChanges(diff []difflib.DiffRecord) int {
	changes := 0
	for _, record := range diff {
		if record.Delta != difflib.Common {
			changes++
		}
	}
	return changes
}

func doDiff(key string, oldContent *manifest.MappingResult, newContent *manifest.MappingResult, suppressedKinds []string, showSecrets bool, stripTrailingCR bool, context int) {
	if oldContent != nil && newContent != nil && oldContent.Content == newContent.Content {
		return
	}

	if !showSecrets {
		redactSecrets(oldContent, newContent)
	}

	if oldContent == nil {
		emptyMapping := &manifest.MappingResult{}
		diffs := diffMappingResults(emptyMapping, newContent, stripTrailingCR)
		report.addEntry(key, suppressedKinds, newContent.Kind, context, diffs, "ADD")
	} else if newContent == nil {
		emptyMapping := &manifest.MappingResult{}
		diffs := diffMappingResults(oldContent, emptyMapping, stripTrailingCR)
		report.addEntry(key, suppressedKinds, oldContent.Kind, context, diffs, "REMOVE")
	} else {
		diffs := diffMappingResults(oldContent, newContent, stripTrailingCR)
		if actualChanges(diffs) > 0 {
			report.addEntry(key, suppressedKinds, oldContent.Kind, context, diffs, "MODIFY")
		}
	}
}

func redactSecrets(old, new *manifest.MappingResult) {
	if (old != nil && old.Kind != "Secret") || (new != nil && new.Kind != "Secret") {
		return
	}
	serializer := json.NewYAMLSerializer(json.DefaultMetaFactory, scheme.Scheme,
		scheme.Scheme)
	var oldSecret, newSecret v1.Secret

	if old != nil {
		if err := yaml.NewYAMLToJSONDecoder(bytes.NewBufferString(old.Content)).Decode(&oldSecret); err != nil {
			old.Content = fmt.Sprintf("Error parsing old secret: %s", err)
		}
	}
	if new != nil {
		if err := yaml.NewYAMLToJSONDecoder(bytes.NewBufferString(new.Content)).Decode(&newSecret); err != nil {
			new.Content = fmt.Sprintf("Error parsing new secret: %s", err)
		}
	}
	if old != nil {
		oldSecret.StringData = make(map[string]string, len(oldSecret.Data))
		for k, v := range oldSecret.Data {
			if new != nil && bytes.Equal(v, newSecret.Data[k]) {
				oldSecret.StringData[k] = fmt.Sprintf("REDACTED # (%d bytes)", len(v))
			} else {
				oldSecret.StringData[k] = fmt.Sprintf("-------- # (%d bytes)", len(v))
			}
		}
	}
	if new != nil {
		newSecret.StringData = make(map[string]string, len(newSecret.Data))
		for k, v := range newSecret.Data {
			if old != nil && bytes.Equal(v, oldSecret.Data[k]) {
				newSecret.StringData[k] = fmt.Sprintf("REDACTED # (%d bytes)", len(v))
			} else {
				newSecret.StringData[k] = fmt.Sprintf("++++++++ # (%d bytes)", len(v))
			}
		}
	}
	// remove Data field now that we are using StringData for serialization
	var buf bytes.Buffer
	if old != nil {
		oldSecret.Data = nil
		if err := serializer.Encode(&oldSecret, &buf); err != nil {
			new.Content = fmt.Sprintf("Error encoding new secret: %s", err)
		}
		old.Content = getComment(old.Content) + strings.Replace(strings.Replace(buf.String(), "stringData", "data", 1), "  creationTimestamp: null\n", "", 1)
		buf.Reset() //reuse buffer for new secret
	}
	if new != nil {
		newSecret.Data = nil
		if err := serializer.Encode(&newSecret, &buf); err != nil {
			new.Content = fmt.Sprintf("Error encoding new secret: %s", err)
		}
		new.Content = getComment(new.Content) + strings.Replace(strings.Replace(buf.String(), "stringData", "data", 1), "  creationTimestamp: null\n", "", 1)
	}
}

// return the first line of a string if its a comment.
// This gives as the # Source: lines from the rendering
func getComment(s string) string {
	i := strings.Index(s, "\n")
	if i < 0 || !strings.HasPrefix(s, "#") {
		return ""
	}
	return s[:i+1]

}

// Releases reindex the content  based on the template names and pass it to Manifests
func Releases(oldIndex, newIndex map[string]*manifest.MappingResult, suppressedKinds []string, showSecrets bool, context int, output string, stripTrailingCR bool, to io.Writer) bool {
	oldIndex = reIndexForRelease(oldIndex)
	newIndex = reIndexForRelease(newIndex)
	return Manifests(oldIndex, newIndex, suppressedKinds, showSecrets, context, output, stripTrailingCR, to)
}

func diffMappingResults(oldContent *manifest.MappingResult, newContent *manifest.MappingResult, stripTrailingCR bool) []difflib.DiffRecord {
	return diffStrings(oldContent.Content, newContent.Content, stripTrailingCR)
}

func diffStrings(before, after string, stripTrailingCR bool) []difflib.DiffRecord {
	return difflib.Diff(split(before, stripTrailingCR), split(after, stripTrailingCR))
}

func split(value string, stripTrailingCR bool) []string {
	const sep = "\n"
	split := strings.Split(value, sep)
	if !stripTrailingCR {
		return split
	}
	var stripped []string
	for _, s := range split {
		stripped = append(stripped, strings.TrimSuffix(s, "\r"))
	}
	return stripped
}

func printDiffRecords(suppressedKinds []string, kind string, context int, diffs []difflib.DiffRecord, to io.Writer) {
	for _, ckind := range suppressedKinds {

		if ckind == kind {
			str := fmt.Sprintf("+ Changes suppressed on sensitive content of type %s\n", kind)
			fmt.Fprintf(to, ansi.Color(str, "yellow"))
			return
		}
	}

	if context >= 0 {
		distances := calculateDistances(diffs)
		omitting := false
		for i, diff := range diffs {
			if distances[i] > context {
				if !omitting {
					fmt.Fprintln(to, "...")
					omitting = true
				}
			} else {
				omitting = false
				printDiffRecord(diff, to)
			}
		}
	} else {
		for _, diff := range diffs {
			printDiffRecord(diff, to)
		}
	}
}

func printDiffRecord(diff difflib.DiffRecord, to io.Writer) {
	text := diff.Payload

	switch diff.Delta {
	case difflib.RightOnly:
		fmt.Fprintf(to, "%s\n", ansi.Color("+ "+text, "green"))
	case difflib.LeftOnly:
		fmt.Fprintf(to, "%s\n", ansi.Color("- "+text, "red"))
	case difflib.Common:
		if text == "" {
			fmt.Fprintln(to)
		} else {
			fmt.Fprintf(to, "%s\n", "  "+text)
		}
	}
}

// Calculate distance of every diff-line to the closest change
func calculateDistances(diffs []difflib.DiffRecord) map[int]int {
	distances := map[int]int{}

	// Iterate forwards through diffs, set 'distance' based on closest 'change' before this line
	change := -1
	for i, diff := range diffs {
		if diff.Delta != difflib.Common {
			change = i
		}
		distance := math.MaxInt32
		if change != -1 {
			distance = i - change
		}
		distances[i] = distance
	}

	// Iterate backwards through diffs, reduce 'distance' based on closest 'change' after this line
	change = -1
	for i := len(diffs) - 1; i >= 0; i-- {
		diff := diffs[i]
		if diff.Delta != difflib.Common {
			change = i
		}
		if change != -1 {
			distance := change - i
			if distance < distances[i] {
				distances[i] = distance
			}
		}
	}

	return distances
}

// reIndexForRelease based on template names
func reIndexForRelease(index map[string]*manifest.MappingResult) map[string]*manifest.MappingResult {

	// sort the index to iterate map in the same order
	var keys []string
	for key := range index {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	// holds number of object in a single file
	count := make(map[string]int)

	newIndex := make(map[string]*manifest.MappingResult)

	for key := range keys {

		str := strings.Replace(strings.Split(index[keys[key]].Content, "\n")[0], "# Source: ", "", 1)

		if _, ok := newIndex[str]; ok {
			count[str]++
			str += fmt.Sprintf(" %d", count[str])
			newIndex[str] = index[keys[key]]
		} else {
			newIndex[str] = index[keys[key]]
			count[str]++
		}
	}
	return newIndex
}

func sortedKeys(manifests map[string]*manifest.MappingResult) []string {
	var keys []string

	for key := range manifests {
		keys = append(keys, key)
	}

	sort.Strings(keys)

	return keys
}
<|endoftext|>


<commit_before>package mpb_test

import (
	"bytes"
	"sync"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/vbauerster/mpb"
)

func TestBarSetWidth(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	// overwrite default width 80
	customWidth := 60
	bar := p.AddBar(100).SetWidth(customWidth).
		TrimLeftSpace().TrimRightSpace()
	for i := 0; i < 100; i++ {
		bar.Incr(1)
	}
	p.Stop()

	gotWidth := len(buf.Bytes())
	if gotWidth != customWidth+1 { // +1 for new line
		t.Errorf("Expected width: %d, got: %d\n", customWidth, gotWidth)
	}
}

func TestBarSetInvalidWidth(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	bar := p.AddBar(100).SetWidth(1).
		TrimLeftSpace().TrimRightSpace()
	for i := 0; i < 100; i++ {
		bar.Incr(1)
	}
	p.Stop()

	wantWidth := 80
	gotWidth := len(buf.Bytes())
	if gotWidth != wantWidth+1 { // +1 for new line
		t.Errorf("Expected width: %d, got: %d\n", wantWidth, gotWidth)
	}
}

func TestBarFormat(t *testing.T) {
	var buf bytes.Buffer
	cancel := make(chan struct{})
	p := mpb.New().WithCancel(cancel).SetOut(&buf)
	customFormat := "(#>_)"
	bar := p.AddBar(100).Format(customFormat).
		TrimLeftSpace().TrimRightSpace()

	go func() {
		for i := 0; i < 100; i++ {
			time.Sleep(10 * time.Millisecond)
			bar.Incr(1)
		}
	}()

	time.Sleep(250 * time.Millisecond)
	close(cancel)
	p.Stop()

	// removing new line
	bytes := removeLastRune(buf.Bytes())

	seen := make(map[rune]bool)
	for _, r := range string(bytes) {
		if !seen[r] {
			seen[r] = true
		}
	}
	for _, r := range customFormat {
		if !seen[r] {
			t.Errorf("Rune %#U not found in bar\n", r)
		}
	}
}

func TestBarInvalidFormat(t *testing.T) {
	var buf bytes.Buffer
	customWidth := 60
	p := mpb.New().SetWidth(customWidth).SetOut(&buf)
	customFormat := "(#>=_)"
	bar := p.AddBar(100).Format(customFormat).
		TrimLeftSpace().TrimRightSpace()

	for i := 0; i < 100; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	bytes := removeLastRune(buf.Bytes())
	got := string(bytes[len(bytes)-customWidth:])
	want := "[==========================================================]"
	if got != want {
		t.Errorf("Expected format: %s, got %s\n", want, got)
	}
}

func TestBarInProgress(t *testing.T) {
	var buf bytes.Buffer
	cancel := make(chan struct{})
	p := mpb.New().WithCancel(cancel).SetOut(&buf)
	bar := p.AddBar(100).TrimLeftSpace().TrimRightSpace()

	stopped := make(chan struct{})

	go func() {
		defer close(stopped)
		for bar.InProgress() {
			time.Sleep(10 * time.Millisecond)
			bar.Incr(1)
		}
	}()

	time.Sleep(250 * time.Millisecond)
	close(cancel)
	p.Stop()

	select {
	case <-stopped:
	case <-time.After(300 * time.Millisecond):
		t.Error("bar.InProgress returns true after cancel")
	}
}

func TestGetSpinner(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	bar := p.AddBar(0).TrimLeftSpace().TrimRightSpace()

	for i := 0; i < 100; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	spinnerChars := []byte(`-\|/`)
	seen := make(map[byte]bool)
	for _, b := range buf.Bytes() {
		if !seen[b] {
			seen[b] = true
		}
	}
	for _, b := range spinnerChars {
		if !seen[b] {
			t.Errorf("Char %#U not found in bar's output\n", b)
		}
	}
}

func TestBarGetID(t *testing.T) {
	var wg sync.WaitGroup
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)

	numBars := 3
	wg.Add(numBars)

	bars := make([]*mpb.Bar, numBars)
	for i := 0; i < numBars; i++ {
		bars[i] = p.AddBarWithID(i, 100)

		go func(bar *mpb.Bar) {
			defer wg.Done()
			for i := 0; i < 100; i++ {
				time.Sleep(10 * time.Millisecond)
				bar.Incr(1)
			}
		}(bars[i])
	}

	for wantID, bar := range bars {
		gotID := bar.GetID()
		if gotID != wantID {
			t.Errorf("Expected bar id: %d, got %d\n", wantID, gotID)
		}
	}

	wg.Wait()
	p.Stop()
}

func removeLastRune(bytes []byte) []byte {
	_, size := utf8.DecodeLastRune(bytes)
	return bytes[:len(bytes)-size]
}
<commit_msg>TestBarIncrWithReFill<commit_after>package mpb_test

import (
	"bytes"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"
	"unicode/utf8"

	"github.com/vbauerster/mpb"
)

func TestBarSetWidth(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	// overwrite default width 80
	customWidth := 60
	bar := p.AddBar(100).SetWidth(customWidth).
		TrimLeftSpace().TrimRightSpace()
	for i := 0; i < 100; i++ {
		bar.Incr(1)
	}
	p.Stop()

	gotWidth := len(buf.Bytes())
	if gotWidth != customWidth+1 { // +1 for new line
		t.Errorf("Expected width: %d, got: %d\n", customWidth, gotWidth)
	}
}

func TestBarSetInvalidWidth(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	bar := p.AddBar(100).SetWidth(1).
		TrimLeftSpace().TrimRightSpace()
	for i := 0; i < 100; i++ {
		bar.Incr(1)
	}
	p.Stop()

	wantWidth := 80
	gotWidth := len(buf.Bytes())
	if gotWidth != wantWidth+1 { // +1 for new line
		t.Errorf("Expected width: %d, got: %d\n", wantWidth, gotWidth)
	}
}

func TestBarFormat(t *testing.T) {
	var buf bytes.Buffer
	cancel := make(chan struct{})
	p := mpb.New().WithCancel(cancel).SetOut(&buf)
	customFormat := "(#>_)"
	bar := p.AddBar(100).Format(customFormat).
		TrimLeftSpace().TrimRightSpace()

	go func() {
		for i := 0; i < 100; i++ {
			time.Sleep(10 * time.Millisecond)
			bar.Incr(1)
		}
	}()

	time.Sleep(250 * time.Millisecond)
	close(cancel)
	p.Stop()

	// removing new line
	bytes := removeLastRune(buf.Bytes())

	seen := make(map[rune]bool)
	for _, r := range string(bytes) {
		if !seen[r] {
			seen[r] = true
		}
	}
	for _, r := range customFormat {
		if !seen[r] {
			t.Errorf("Rune %#U not found in bar\n", r)
		}
	}
}

func TestBarInvalidFormat(t *testing.T) {
	var buf bytes.Buffer
	customWidth := 60
	p := mpb.New().SetWidth(customWidth).SetOut(&buf)
	customFormat := "(#>=_)"
	bar := p.AddBar(100).Format(customFormat).
		TrimLeftSpace().TrimRightSpace()

	for i := 0; i < 100; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	bytes := removeLastRune(buf.Bytes())
	got := string(bytes[len(bytes)-customWidth:])
	want := "[==========================================================]"
	if got != want {
		t.Errorf("Expected format: %s, got %s\n", want, got)
	}
}

func TestBarInProgress(t *testing.T) {
	var buf bytes.Buffer
	cancel := make(chan struct{})
	p := mpb.New().WithCancel(cancel).SetOut(&buf)
	bar := p.AddBar(100).TrimLeftSpace().TrimRightSpace()

	stopped := make(chan struct{})

	go func() {
		defer close(stopped)
		for bar.InProgress() {
			time.Sleep(10 * time.Millisecond)
			bar.Incr(1)
		}
	}()

	time.Sleep(250 * time.Millisecond)
	close(cancel)
	p.Stop()

	select {
	case <-stopped:
	case <-time.After(300 * time.Millisecond):
		t.Error("bar.InProgress returns true after cancel")
	}
}

func TestGetSpinner(t *testing.T) {
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)
	bar := p.AddBar(0).TrimLeftSpace().TrimRightSpace()

	for i := 0; i < 100; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	spinnerChars := []byte(`-\|/`)
	seen := make(map[byte]bool)
	for _, b := range buf.Bytes() {
		if !seen[b] {
			seen[b] = true
		}
	}
	for _, b := range spinnerChars {
		if !seen[b] {
			t.Errorf("Char %#U not found in bar's output\n", b)
		}
	}
}

func TestBarGetID(t *testing.T) {
	var wg sync.WaitGroup
	var buf bytes.Buffer
	p := mpb.New().SetOut(&buf)

	numBars := 3
	wg.Add(numBars)

	bars := make([]*mpb.Bar, numBars)
	for i := 0; i < numBars; i++ {
		bars[i] = p.AddBarWithID(i, 100)

		go func(bar *mpb.Bar) {
			defer wg.Done()
			for i := 0; i < 100; i++ {
				time.Sleep(10 * time.Millisecond)
				bar.Incr(1)
			}
		}(bars[i])
	}

	for wantID, bar := range bars {
		gotID := bar.GetID()
		if gotID != wantID {
			t.Errorf("Expected bar id: %d, got %d\n", wantID, gotID)
		}
	}

	wg.Wait()
	p.Stop()
}

func TestBarIncrWithReFill(t *testing.T) {
	var buf bytes.Buffer

	width := 100
	p := mpb.New().SetWidth(width).SetOut(&buf)

	total := 100
	refill := 30
	delta := total - refill
	refillChar := '+'
	bar := p.AddBar(int64(total)).TrimLeftSpace().TrimRightSpace()

	bar.IncrWithReFill(refill, &mpb.Refill{Char: refillChar})

	for i := 0; i < delta; i++ {
		time.Sleep(10 * time.Millisecond)
		bar.Incr(1)
	}

	p.Stop()

	bytes := removeLastRune(buf.Bytes())

	gotBar := string(bytes[len(bytes)-width:])
	wantBar := fmt.Sprintf("[%s%s]",
		strings.Repeat(string(refillChar), refill-1),
		strings.Repeat("=", delta-1))
	if gotBar != wantBar {
		t.Errorf("Want bar: %s, got bar: %s\n", wantBar, gotBar)
	}
}

func removeLastRune(bytes []byte) []byte {
	_, size := utf8.DecodeLastRune(bytes)
	return bytes[:len(bytes)-size]
}
<|endoftext|>