<commit_before>package mysql

import (
	"database/sql"
	"fmt"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/solefaucet/sole-server/errors"
	"github.com/solefaucet/sole-server/models"
)

// GetRewardIncomes get user's reward incomes
func (s Storage) GetRewardIncomes(userID int64, limit, offset int64) ([]models.Income, error) {
	rawSQL := "SELECT * FROM incomes WHERE `user_id` = ? AND `type` = ? ORDER BY `id` DESC LIMIT ? OFFSET ?"
	args := []interface{}{userID, models.IncomeTypeReward, limit, offset}
	incomes := []models.Income{}
	err := s.selects(&incomes, rawSQL, args...)
	return incomes, err
}

// GetNumberOfRewardIncomes gets number of user's reward incomes
func (s Storage) GetNumberOfRewardIncomes(userID int64) (int64, error) {
	var count int64
	err := s.db.QueryRowx("SELECT COUNT(*) FROM incomes WHERE `user_id` = ? AND `type` = ?", userID, models.IncomeTypeReward).Scan(&count)
	return count, err
}

// CreateRewardIncome creates a new reward type income
func (s Storage) CreateRewardIncome(income models.Income, now time.Time) error {
	tx := s.db.MustBegin()

	if err := createRewardIncomeWithTx(tx, income, now); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create reward income commit transaction error: %v", err)
	}

	return nil
}

func createRewardIncomeWithTx(tx *sqlx.Tx, income models.Income, now time.Time) error {
	totalReward := income.Income

	// insert income into incomes table
	if _, err := addIncome(tx, income); err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err := incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update user rewarded_at
	if _, err := tx.Exec("UPDATE users SET `rewarded_at` = ? WHERE `id` = ?", now, income.UserID); err != nil {
		return err
	}

	// update referer balance
	if rowAffected, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	} else if rowAffected == 1 {
		totalReward += income.RefererIncome
	}

	// update total reward
	if err := incrementTotalReward(tx, totalReward, now); err != nil {
		return err
	}

	return nil
}

// insert reward income into incomes table
func addIncome(tx *sqlx.Tx, income models.Income) (sql.Result, error) {
	result, err := tx.NamedExec("INSERT INTO incomes (`user_id`, `referer_id`, `type`, `income`, `referer_income`) VALUES (:user_id, :referer_id, :type, :income, :referer_income)", income)
	if err != nil {
		return nil, fmt.Errorf("add income error: %v", err)
	}

	return result, nil
}

// increment user balance, total_income, referer_total_income
func incrementUserBalance(tx *sqlx.Tx, userID int64, delta, refererDelta float64) error {
	rawSQL := "UPDATE users SET `balance` = `balance` + ?, `total_income` = `total_income` + ?, `referer_total_income` = `referer_total_income` + ? WHERE id = ?"
	args := []interface{}{delta, delta, refererDelta, userID}
	if result, err := tx.Exec(rawSQL, args...); err != nil {
		return fmt.Errorf("increment user balance error: %v", err)
	} else if rowAffected, _ := result.RowsAffected(); rowAffected != 1 {
		return fmt.Errorf("increment user balance affected %v rows", rowAffected)
	}

	return nil
}

// increment referer balance
func incrementRefererBalance(tx *sqlx.Tx, refererID int64, delta float64) (int64, error) {
	result, err := tx.NamedExec("UPDATE users SET `balance` = `balance` + :delta, `total_income_from_referees` = `total_income_from_referees` + :delta WHERE id = :id", map[string]interface{}{
		"id":    refererID,
		"delta": delta,
	})
	if err != nil {
		return 0, fmt.Errorf("increment referer balance error: %v", err)
	}

	rowAffected, _ := result.RowsAffected()
	return rowAffected, nil
}

// increment total reward
func incrementTotalReward(tx *sqlx.Tx, totalReward float64, now time.Time) error {
	sql := "INSERT INTO total_rewards (`total`, `created_at`) VALUES (:delta, :created_at) ON DUPLICATE KEY UPDATE `total` = `total` + :delta"
	args := map[string]interface{}{
		"delta":      totalReward,
		"created_at": now,
	}

	if _, err := tx.NamedExec(sql, args); err != nil {
		return fmt.Errorf("increment total reward error: %v", err)
	}

	return nil
}

// GetOfferwowEventByID finds offerwow event by event id
func (s Storage) GetOfferwowEventByID(eventID string) (models.OfferwowEvent, error) {
	event := models.OfferwowEvent{}
	err := s.db.Get(&event, "SELECT * FROM `offerwow` WHERE `event_id` = ?", eventID)

	if err != nil {
		if err == sql.ErrNoRows {
			return event, errors.ErrNotFound
		}

		return event, fmt.Errorf("query offerwow event by id error: %v", err)
	}

	return event, nil
}

// CreateOfferwowIncome creates a new offerwow type income
func (s Storage) CreateOfferwowIncome(income models.Income, eventID string) error {
	tx := s.db.MustBegin()

	if err := createOfferwowIncomeWithTx(tx, income, eventID); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create reward income commit transaction error: %v", err)
	}

	return nil
}

func createOfferwowIncomeWithTx(tx *sqlx.Tx, income models.Income, eventID string) error {
	// update user balance, total_income, referer_total_income
	if err := incrementUserBalanceByOfferwowIncome(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update referer balance
	if _, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	}

	// insert offerwow income into incomes table
	result, err := addIncome(tx, income)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// insert offerwow event
	offerwowEvent := models.OfferwowEvent{
		EventID:  eventID,
		IncomeID: id,
		Amount:   income.Income,
	}
	_, err = tx.NamedExec("INSERT INTO `offerwow` (`event_id`, `income_id`, `amount`) VALUE (:event_id, :income_id, :amount)", offerwowEvent)
	return err
}

// update user balance, total_income, referer_total_income
func incrementUserBalanceByOfferwowIncome(tx *sqlx.Tx, userID int64, delta, refererDelta float64) error {
	rawSQL := "UPDATE users SET `balance` = `balance` + ?, `total_income` = `total_income` + ?, `referer_total_income` = `referer_total_income` + ? WHERE id = ?"
	args := []interface{}{delta, delta, refererDelta, userID}
	if result, err := tx.Exec(rawSQL, args...); err != nil {
		return fmt.Errorf("create offerwow income update user balance error: %v", err)
	} else if rowAffected, _ := result.RowsAffected(); rowAffected != 1 {
		return fmt.Errorf("create offerwow income update user balance affected %v rows", rowAffected)
	}

	return nil
}

// GetSuperrewardsOfferByID finds superrewards offer by transaction Id and user id
func (s Storage) GetSuperrewardsOfferByID(transactionID string, userID int64) (models.SuperrewardsOffer, error) {
	offer := models.SuperrewardsOffer{}
	err := s.db.Get(&offer, "SELECT * FROM `superrewards` WHERE `transaction_id` = ? AND `user_id` = ?", transactionID, userID)

	if err != nil {
		if err == sql.ErrNoRows {
			return offer, errors.ErrNotFound
		}

		return offer, fmt.Errorf("query superrewards offer by id error: %v", err)
	}

	return offer, nil
}

// CreateSuperrewardsIncome creates a new superrewards type income
func (s Storage) CreateSuperrewardsIncome(income models.Income, transactionID, offerID string) error {
	tx := s.db.MustBegin()

	if err := createSuperrewardsIncomeWithTx(tx, income, transactionID, offerID); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create reward income commit transaction error: %v", err)
	}

	return nil
}

func createSuperrewardsIncomeWithTx(tx *sqlx.Tx, income models.Income, transactionID, offerID string) error {
	// update user balance, total_income, referer_total_income
	if err := incrementUserBalanceBySuperrewardsIncome(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update referer balance
	if _, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	}

	// insert superrewards income into incomes table
	result, err := addIncome(tx, income)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// insert superrewards offer
	offer := models.SuperrewardsOffer{
		IncomeID:      id,
		UserID:        income.UserID,
		TransactionID: transactionID,
		OfferID:       offerID,
		Amount:        income.Income,
	}
	_, err = tx.NamedExec("INSERT INTO `superrewards` (`income_id`, `user_id`, `transaction_id`, `offer_id`, `amount`) VALUE (:income_id, :user_id, :transaction_id, :offer_id, :amount)", offer)
	return err
}

// update user balance, total_income, referer_total_income
func incrementUserBalanceBySuperrewardsIncome(tx *sqlx.Tx, userID int64, delta, refererDelta float64) error {
	rawSQL := "UPDATE users SET `balance` = `balance` + ?, `total_income` = `total_income` + ?, `referer_total_income` = `referer_total_income` + ? WHERE id = ?"
	args := []interface{}{delta, delta, refererDelta, userID}
	if result, err := tx.Exec(rawSQL, args...); err != nil {
		return fmt.Errorf("create superrewards income update user balance error: %v", err)
	} else if rowAffected, _ := result.RowsAffected(); rowAffected != 1 {
		return fmt.Errorf("create superrewards income update user balance affected %v rows", rowAffected)
	}

	return nil
}
<commit_msg>Refactor createSuperrewardsIncomeWithTx<commit_after>package mysql

import (
	"database/sql"
	"fmt"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/solefaucet/sole-server/errors"
	"github.com/solefaucet/sole-server/models"
)

// GetRewardIncomes get user's reward incomes
func (s Storage) GetRewardIncomes(userID int64, limit, offset int64) ([]models.Income, error) {
	rawSQL := "SELECT * FROM incomes WHERE `user_id` = ? AND `type` = ? ORDER BY `id` DESC LIMIT ? OFFSET ?"
	args := []interface{}{userID, models.IncomeTypeReward, limit, offset}
	incomes := []models.Income{}
	err := s.selects(&incomes, rawSQL, args...)
	return incomes, err
}

// GetNumberOfRewardIncomes gets number of user's reward incomes
func (s Storage) GetNumberOfRewardIncomes(userID int64) (int64, error) {
	var count int64
	err := s.db.QueryRowx("SELECT COUNT(*) FROM incomes WHERE `user_id` = ? AND `type` = ?", userID, models.IncomeTypeReward).Scan(&count)
	return count, err
}

// CreateRewardIncome creates a new reward type income
func (s Storage) CreateRewardIncome(income models.Income, now time.Time) error {
	tx := s.db.MustBegin()

	if err := createRewardIncomeWithTx(tx, income, now); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create reward income commit transaction error: %v", err)
	}

	return nil
}

func createRewardIncomeWithTx(tx *sqlx.Tx, income models.Income, now time.Time) error {
	totalReward := income.Income

	// insert income into incomes table
	if _, err := addIncome(tx, income); err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err := incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update user rewarded_at
	if _, err := tx.Exec("UPDATE users SET `rewarded_at` = ? WHERE `id` = ?", now, income.UserID); err != nil {
		return err
	}

	// update referer balance
	if rowAffected, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	} else if rowAffected == 1 {
		totalReward += income.RefererIncome
	}

	// update total reward
	if err := incrementTotalReward(tx, totalReward, now); err != nil {
		return err
	}

	return nil
}

// insert reward income into incomes table
func addIncome(tx *sqlx.Tx, income models.Income) (sql.Result, error) {
	result, err := tx.NamedExec("INSERT INTO incomes (`user_id`, `referer_id`, `type`, `income`, `referer_income`) VALUES (:user_id, :referer_id, :type, :income, :referer_income)", income)
	if err != nil {
		return nil, fmt.Errorf("add income error: %v", err)
	}

	return result, nil
}

// increment user balance, total_income, referer_total_income
func incrementUserBalance(tx *sqlx.Tx, userID int64, delta, refererDelta float64) error {
	rawSQL := "UPDATE users SET `balance` = `balance` + ?, `total_income` = `total_income` + ?, `referer_total_income` = `referer_total_income` + ? WHERE id = ?"
	args := []interface{}{delta, delta, refererDelta, userID}
	if result, err := tx.Exec(rawSQL, args...); err != nil {
		return fmt.Errorf("increment user balance error: %v", err)
	} else if rowAffected, _ := result.RowsAffected(); rowAffected != 1 {
		return fmt.Errorf("increment user balance affected %v rows", rowAffected)
	}

	return nil
}

// increment referer balance
func incrementRefererBalance(tx *sqlx.Tx, refererID int64, delta float64) (int64, error) {
	result, err := tx.NamedExec("UPDATE users SET `balance` = `balance` + :delta, `total_income_from_referees` = `total_income_from_referees` + :delta WHERE id = :id", map[string]interface{}{
		"id":    refererID,
		"delta": delta,
	})
	if err != nil {
		return 0, fmt.Errorf("increment referer balance error: %v", err)
	}

	rowAffected, _ := result.RowsAffected()
	return rowAffected, nil
}

// increment total reward
func incrementTotalReward(tx *sqlx.Tx, totalReward float64, now time.Time) error {
	sql := "INSERT INTO total_rewards (`total`, `created_at`) VALUES (:delta, :created_at) ON DUPLICATE KEY UPDATE `total` = `total` + :delta"
	args := map[string]interface{}{
		"delta":      totalReward,
		"created_at": now,
	}

	if _, err := tx.NamedExec(sql, args); err != nil {
		return fmt.Errorf("increment total reward error: %v", err)
	}

	return nil
}

// GetOfferwowEventByID finds offerwow event by event id
func (s Storage) GetOfferwowEventByID(eventID string) (models.OfferwowEvent, error) {
	event := models.OfferwowEvent{}
	err := s.db.Get(&event, "SELECT * FROM `offerwow` WHERE `event_id` = ?", eventID)

	if err != nil {
		if err == sql.ErrNoRows {
			return event, errors.ErrNotFound
		}

		return event, fmt.Errorf("query offerwow event by id error: %v", err)
	}

	return event, nil
}

// CreateOfferwowIncome creates a new offerwow type income
func (s Storage) CreateOfferwowIncome(income models.Income, eventID string) error {
	tx := s.db.MustBegin()

	if err := createOfferwowIncomeWithTx(tx, income, eventID); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create reward income commit transaction error: %v", err)
	}

	return nil
}

func createOfferwowIncomeWithTx(tx *sqlx.Tx, income models.Income, eventID string) error {
	// update user balance, total_income, referer_total_income
	if err := incrementUserBalanceByOfferwowIncome(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update referer balance
	if _, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	}

	// insert offerwow income into incomes table
	result, err := addIncome(tx, income)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// insert offerwow event
	offerwowEvent := models.OfferwowEvent{
		EventID:  eventID,
		IncomeID: id,
		Amount:   income.Income,
	}
	_, err = tx.NamedExec("INSERT INTO `offerwow` (`event_id`, `income_id`, `amount`) VALUE (:event_id, :income_id, :amount)", offerwowEvent)
	return err
}

// update user balance, total_income, referer_total_income
func incrementUserBalanceByOfferwowIncome(tx *sqlx.Tx, userID int64, delta, refererDelta float64) error {
	rawSQL := "UPDATE users SET `balance` = `balance` + ?, `total_income` = `total_income` + ?, `referer_total_income` = `referer_total_income` + ? WHERE id = ?"
	args := []interface{}{delta, delta, refererDelta, userID}
	if result, err := tx.Exec(rawSQL, args...); err != nil {
		return fmt.Errorf("create offerwow income update user balance error: %v", err)
	} else if rowAffected, _ := result.RowsAffected(); rowAffected != 1 {
		return fmt.Errorf("create offerwow income update user balance affected %v rows", rowAffected)
	}

	return nil
}

// GetSuperrewardsOfferByID finds superrewards offer by transaction Id and user id
func (s Storage) GetSuperrewardsOfferByID(transactionID string, userID int64) (models.SuperrewardsOffer, error) {
	offer := models.SuperrewardsOffer{}
	err := s.db.Get(&offer, "SELECT * FROM `superrewards` WHERE `transaction_id` = ? AND `user_id` = ?", transactionID, userID)

	if err != nil {
		if err == sql.ErrNoRows {
			return offer, errors.ErrNotFound
		}

		return offer, fmt.Errorf("query superrewards offer by id error: %v", err)
	}

	return offer, nil
}

// CreateSuperrewardsIncome creates a new superrewards type income
func (s Storage) CreateSuperrewardsIncome(income models.Income, transactionID, offerID string) error {
	tx := s.db.MustBegin()

	if err := createSuperrewardsIncomeWithTx(tx, income, transactionID, offerID); err != nil {
		tx.Rollback()
		return err
	}

	// commit
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("create reward income commit transaction error: %v", err)
	}

	return nil
}

func createSuperrewardsIncomeWithTx(tx *sqlx.Tx, income models.Income, transactionID, offerID string) error {
	// insert superrewards income into incomes table
	result, err := addIncome(tx, income)
	if err != nil {
		return err
	}
	id, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// update user balance, total_income, referer_total_income
	if err := incrementUserBalance(tx, income.UserID, income.Income, income.RefererIncome); err != nil {
		return err
	}

	// update referer balance
	if _, err := incrementRefererBalance(tx, income.RefererID, income.RefererIncome); err != nil {
		return err
	}

	// insert superrewards offer
	offer := models.SuperrewardsOffer{
		IncomeID:      id,
		UserID:        income.UserID,
		TransactionID: transactionID,
		OfferID:       offerID,
		Amount:        income.Income,
	}
	_, err = tx.NamedExec("INSERT INTO `superrewards` (`income_id`, `user_id`, `transaction_id`, `offer_id`, `amount`) VALUE (:income_id, :user_id, :transaction_id, :offer_id, :amount)", offer)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"

	"github.com/bazelbuild/rules_typescript/devserver/concatjs"
	"github.com/bazelbuild/rules_typescript/devserver/devserver"
	"github.com/bazelbuild/rules_typescript/devserver/runfiles"
)

var (
	port = flag.Int("port", 5432, "server port to listen on")
	// The "base" CLI flag is only kept because within Google3 because removing would be a breaking change due to
	// ConcatJS and "devserver/devserver.go" still respecting the specified base flag.
	base            = flag.String("base", "", "server base (required, runfiles of the binary)")
	pkgs            = flag.String("packages", "", "root package(s) to serve, comma-separated")
	manifest        = flag.String("manifest", "", "sources manifest (.MF)")
	scriptsManifest = flag.String("scripts_manifest", "", "preScripts manifest (.MF)")
	servingPath     = flag.String("serving_path", "/_/ts_scripts.js", "path to serve the combined sources at")
	entryModule     = flag.String("entry_module", "", "entry module name")
)

func main() {
	flag.Parse()

	if len(*pkgs) == 0 || (*manifest == "") || (*scriptsManifest == "") {
		fmt.Fprintf(os.Stderr, "Required argument not set\n")
		os.Exit(1)
	}

	manifestPath, err := runfiles.Runfile(*base, *scriptsManifest)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to find scripts_manifest in runfiles: %v\n", err)
		os.Exit(1)
	}

	scriptFiles, err := manifestFiles(manifestPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read scripts_manifest: %v\n", err)
		os.Exit(1)
	}

	if !strings.HasPrefix(*servingPath, "/") {
		fmt.Fprintf(os.Stderr, "The specified serving_path does not start with a slash. "+
			"This causes the serving path to not have any effect.\n")
		os.Exit(1)
	}

	preScripts := make([]string, 0, 100)
	postScripts := make([]string, 0, 1)

	// Include the livereload script if IBAZEL_LIVERELOAD_URL is set.
	livereloadUrl := os.Getenv("IBAZEL_LIVERELOAD_URL")
	if livereloadUrl != "" {
		fmt.Printf("Serving livereload script from %s\n", livereloadUrl)
		livereloadLoaderSnippet := fmt.Sprintf(`(function(){
	const script = document.createElement('script');
	script.src = "%s";
	document.head.appendChild(script);
})();`, livereloadUrl)
		preScripts = append(preScripts, livereloadLoaderSnippet)
	}

	// Include the profiler script if IBAZEL_PROFILER_URL is set.
	profilerScriptURL := os.Getenv("IBAZEL_PROFILER_URL")
	if profilerScriptURL != "" {
		fmt.Printf("Serving profiler script from %s\n", profilerScriptURL)
		profilerLoaderSnippet := fmt.Sprintf(`(function(){
	const script = document.createElement('script');
	script.src = "%s";
	document.head.appendChild(script);
})();`, profilerScriptURL)
		preScripts = append(preScripts, profilerLoaderSnippet)
	}

	// Include all user scripts in preScripts. This should always include
	// the requirejs script which is added to scriptFiles by the devserver
	// skylark rule.
	for _, v := range scriptFiles {
		runfile, err := runfiles.Runfile(*base, v)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Could not find runfile %s, got error %s", v, err)
		}

		js, err := loadScript(runfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to read script %s: %v\n", v, err)
		} else {
			preScripts = append(preScripts, js)
		}
	}

	// If the entryModule is set then add a snippet to load
	// the application to postScripts to be outputted after the sources
	if *entryModule != "" {
		postScripts = append(postScripts, fmt.Sprintf("require([\"%s\"]);", *entryModule))
	}

	http.Handle(*servingPath, concatjs.ServeConcatenatedJS(*manifest, *base, preScripts, postScripts,
		&RunfileFileSystem{}))
	pkgList := strings.Split(*pkgs, ",")
	http.HandleFunc("/", devserver.CreateFileHandler(*servingPath, *manifest, pkgList, *base))

	h, err := os.Hostname()
	if err != nil {
		h = "localhost"
	}
	// Detect if we are running in a linux container inside ChromeOS
	// If so, we assume you want to use the native browser (outside the container)
	// so you'll need to modify the hostname to access the server
	if _, err := os.Stat("/etc/apt/sources.list.d/cros.list"); err == nil {
		h = h + ".linux.test"
	}

	fmt.Printf("Server listening on http://%s:%d/\n", h, *port)
	fmt.Fprintln(os.Stderr, http.ListenAndServe(fmt.Sprintf(":%d", *port), nil).Error())
	os.Exit(1)
}

func loadScript(path string) (string, error) {
	buf, err := ioutil.ReadFile(path)
	if err != nil {
		return "", err
	}
	return string(buf), nil
}

// manifestFiles parses a manifest, returning a list of the files in the manifest.
func manifestFiles(manifest string) ([]string, error) {
	f, err := os.Open(manifest)
	if err != nil {
		return nil, fmt.Errorf("could not read manifest %s: %s", manifest, err)
	}
	defer f.Close()
	return manifestFilesFromReader(f)
}

// manifestFilesFromReader is a helper for manifestFiles, split out for testing.
func manifestFilesFromReader(r io.Reader) ([]string, error) {
	var lines []string
	s := bufio.NewScanner(r)
	for s.Scan() {
		path := s.Text()
		if path == "" {
			continue
		}
		lines = append(lines, path)
	}
	if err := s.Err(); err != nil {
		return nil, err
	}

	return lines, nil
}
<commit_msg>script names are prepended to scripts that are eval'ed but the names are missing for prescripts.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"

	"github.com/bazelbuild/rules_typescript/devserver/concatjs"
	"github.com/bazelbuild/rules_typescript/devserver/devserver"
	"github.com/bazelbuild/rules_typescript/devserver/runfiles"
)

var (
	port = flag.Int("port", 5432, "server port to listen on")
	// The "base" CLI flag is only kept because within Google3 because removing would be a breaking change due to
	// ConcatJS and "devserver/devserver.go" still respecting the specified base flag.
	base            = flag.String("base", "", "server base (required, runfiles of the binary)")
	pkgs            = flag.String("packages", "", "root package(s) to serve, comma-separated")
	manifest        = flag.String("manifest", "", "sources manifest (.MF)")
	scriptsManifest = flag.String("scripts_manifest", "", "preScripts manifest (.MF)")
	servingPath     = flag.String("serving_path", "/_/ts_scripts.js", "path to serve the combined sources at")
	entryModule     = flag.String("entry_module", "", "entry module name")
)

func main() {
	flag.Parse()

	if len(*pkgs) == 0 || (*manifest == "") || (*scriptsManifest == "") {
		fmt.Fprintf(os.Stderr, "Required argument not set\n")
		os.Exit(1)
	}

	manifestPath, err := runfiles.Runfile(*base, *scriptsManifest)

	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to find scripts_manifest in runfiles: %v\n", err)
		os.Exit(1)
	}

	scriptFiles, err := manifestFiles(manifestPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to read scripts_manifest: %v\n", err)
		os.Exit(1)
	}

	if !strings.HasPrefix(*servingPath, "/") {
		fmt.Fprintf(os.Stderr, "The specified serving_path does not start with a slash. "+
			"This causes the serving path to not have any effect.\n")
		os.Exit(1)
	}

	preScripts := make([]string, 0, 100)
	postScripts := make([]string, 0, 1)

	// Include the livereload script if IBAZEL_LIVERELOAD_URL is set.
	livereloadUrl := os.Getenv("IBAZEL_LIVERELOAD_URL")
	if livereloadUrl != "" {
		fmt.Printf("Serving livereload script from %s\n", livereloadUrl)
		livereloadLoaderSnippet := fmt.Sprintf(`(function(){
	const script = document.createElement('script');
	script.src = "%s";
	document.head.appendChild(script);
})();`, livereloadUrl)
		preScripts = append(preScripts, livereloadLoaderSnippet)
	}

	// Include the profiler script if IBAZEL_PROFILER_URL is set.
	profilerScriptURL := os.Getenv("IBAZEL_PROFILER_URL")
	if profilerScriptURL != "" {
		fmt.Printf("Serving profiler script from %s\n", profilerScriptURL)
		profilerLoaderSnippet := fmt.Sprintf(`(function(){
	const script = document.createElement('script');
	script.src = "%s";
	document.head.appendChild(script);
})();`, profilerScriptURL)
		preScripts = append(preScripts, profilerLoaderSnippet)
	}

	// Include all user scripts in preScripts. This should always include
	// the requirejs script which is added to scriptFiles by the devserver
	// skylark rule.
	for _, v := range scriptFiles {
		runfile, err := runfiles.Runfile(*base, v)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Could not find runfile %s, got error %s", v, err)
		}

		js, err := loadScript(runfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Failed to read script %s: %v\n", v, err)
		} else {
			preScripts = append(preScripts, js)
		}
	}

	// If the entryModule is set then add a snippet to load
	// the application to postScripts to be outputted after the sources
	if *entryModule != "" {
		postScripts = append(postScripts, fmt.Sprintf("require([\"%s\"]);", *entryModule))
	}

	http.Handle(*servingPath, concatjs.ServeConcatenatedJS(*manifest, *base, preScripts, postScripts,
		&RunfileFileSystem{}))
	pkgList := strings.Split(*pkgs, ",")
	http.HandleFunc("/", devserver.CreateFileHandler(*servingPath, *manifest, pkgList, *base))

	h, err := os.Hostname()
	if err != nil {
		h = "localhost"
	}
	// Detect if we are running in a linux container inside ChromeOS
	// If so, we assume you want to use the native browser (outside the container)
	// so you'll need to modify the hostname to access the server
	if _, err := os.Stat("/etc/apt/sources.list.d/cros.list"); err == nil {
		h = h + ".linux.test"
	}

	fmt.Printf("Server listening on http://%s:%d/\n", h, *port)
	fmt.Fprintln(os.Stderr, http.ListenAndServe(fmt.Sprintf(":%d", *port), nil).Error())
	os.Exit(1)
}

func loadScript(path string) (string, error) {
	buf, err := ioutil.ReadFile(path)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("// %s\n%s", path, buf), nil
}

// manifestFiles parses a manifest, returning a list of the files in the manifest.
func manifestFiles(manifest string) ([]string, error) {
	f, err := os.Open(manifest)
	if err != nil {
		return nil, fmt.Errorf("could not read manifest %s: %s", manifest, err)
	}
	defer f.Close()
	return manifestFilesFromReader(f)
}

// manifestFilesFromReader is a helper for manifestFiles, split out for testing.
func manifestFilesFromReader(r io.Reader) ([]string, error) {
	var lines []string
	s := bufio.NewScanner(r)
	for s.Scan() {
		path := s.Text()
		if path == "" {
			continue
		}
		lines = append(lines, path)
	}
	if err := s.Err(); err != nil {
		return nil, err
	}

	return lines, nil
}
<|endoftext|>


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

package command

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

	"github.com/coreos/etcd/clientv3"
	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

// NewLeaseRenewerCommand returns the cobra command for "lease-renewer runner".
func NewLeaseRenewerCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "lease-renewer",
		Short: "Performs lease renew operation",
		Run:   runLeaseRenewerFunc,
	}
	return cmd
}

func runLeaseRenewerFunc(cmd *cobra.Command, args []string) {
	if len(args) > 0 {
		ExitWithError(ExitBadArgs, errors.New("lease-renewer does not take any argument"))
	}

	eps := endpointsFromFlag(cmd)
	dialTimeout := dialTimeoutFromCmd(cmd)
	c := newClient(eps, dialTimeout)
	ctx := context.Background()

	for {
		var (
			l   *clientv3.LeaseGrantResponse
			lk  *clientv3.LeaseKeepAliveResponse
			err error
		)
		for {
			l, err = c.Lease.Grant(ctx, 5)
			if err == nil {
				break
			}
		}
		expire := time.Now().Add(time.Duration(l.TTL-1) * time.Second)

		for {
			lk, err = c.Lease.KeepAliveOnce(ctx, l.ID)
			if grpc.Code(err) == codes.NotFound {
				if time.Since(expire) < 0 {
					log.Printf("bad renew! exceeded: %v", time.Since(expire))
					for {
						lk, err = c.Lease.KeepAliveOnce(ctx, l.ID)
						fmt.Println(lk, err)
						time.Sleep(time.Second)
					}
				}
				log.Printf("lost lease %d, expire: %v\n", l.ID, expire)
				break
			}
			if err != nil {
				continue
			}
			expire = time.Now().Add(time.Duration(lk.TTL-1) * time.Second)
			log.Printf("renewed lease %d, expire: %v\n", lk.ID, expire)
			time.Sleep(time.Duration(lk.TTL-2) * time.Second)
		}
	}
}
<commit_msg>etcd-runner: update to use new lease interface<commit_after>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package command

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

	"github.com/coreos/etcd/clientv3"
	"github.com/spf13/cobra"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

// NewLeaseRenewerCommand returns the cobra command for "lease-renewer runner".
func NewLeaseRenewerCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "lease-renewer",
		Short: "Performs lease renew operation",
		Run:   runLeaseRenewerFunc,
	}
	return cmd
}

func runLeaseRenewerFunc(cmd *cobra.Command, args []string) {
	if len(args) > 0 {
		ExitWithError(ExitBadArgs, errors.New("lease-renewer does not take any argument"))
	}

	eps := endpointsFromFlag(cmd)
	dialTimeout := dialTimeoutFromCmd(cmd)
	c := newClient(eps, dialTimeout)
	ctx := context.Background()

	for {
		var (
			l   *clientv3.LeaseGrantResponse
			err error
		)
		for {
			l, err = c.Lease.Grant(ctx, 5)
			if err == nil {
				break
			}
		}
		expire := time.Now().Add(time.Duration(l.TTL-1) * time.Second)

		for {
			lk := c.Lease.KeepAliveOnce(ctx, l.ID)
			err = lk.Err
			if grpc.Code(err) == codes.NotFound {
				if time.Since(expire) < 0 {
					log.Printf("bad renew! exceeded: %v", time.Since(expire))
					for {
						lk = c.Lease.KeepAliveOnce(ctx, l.ID)
						fmt.Println(lk)
						time.Sleep(time.Second)
					}
				}
				log.Printf("lost lease %d, expire: %v\n", l.ID, expire)
				break
			}
			if err != nil {
				continue
			}
			expire = time.Now().Add(time.Duration(lk.TTL-1) * time.Second)
			log.Printf("renewed lease %d, expire: %v\n", lk.ID, expire)
			time.Sleep(time.Duration(lk.TTL-2) * time.Second)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"io"
	"os"
	"syscall"
	"time"
)

// TCPConn is an implementation of the Conn interface for TCP network
// connections.
type TCPConn struct {
	conn
}

func newTCPConn(fd *netFD) *TCPConn {
	return &TCPConn{conn{fd}}
}

// ReadFrom implements the io.ReaderFrom ReadFrom method.
func (c *TCPConn) ReadFrom(r io.Reader) (int64, error) {
	return genericReadFrom(c, r)
}

// CloseRead shuts down the reading side of the TCP connection.
// Most callers should just use Close.
func (c *TCPConn) CloseRead() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.CloseRead()
}

// CloseWrite shuts down the writing side of the TCP connection.
// Most callers should just use Close.
func (c *TCPConn) CloseWrite() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.CloseWrite()
}

// SetLinger sets the behavior of Close() on a connection which still
// has data waiting to be sent or to be acknowledged.
//
// If sec < 0 (the default), Close returns immediately and the
// operating system finishes sending the data in the background.
//
// If sec == 0, Close returns immediately and the operating system
// discards any unsent or unacknowledged data.
//
// If sec > 0, Close blocks for at most sec seconds waiting for data
// to be sent and acknowledged.
func (c *TCPConn) SetLinger(sec int) error {
	return syscall.EPLAN9
}

// SetKeepAlive sets whether the operating system should send
// keepalive messages on the connection.
func (c *TCPConn) SetKeepAlive(keepalive bool) error {
	return syscall.EPLAN9
}

// SetNoDelay controls whether the operating system should delay
// packet transmission in hopes of sending fewer packets (Nagle's
// algorithm).  The default is true (no delay), meaning that data is
// sent as soon as possible after a Write.
func (c *TCPConn) SetNoDelay(noDelay bool) error {
	return syscall.EPLAN9
}

// DialTCP connects to the remote address raddr on the network net,
// which must be "tcp", "tcp4", or "tcp6".  If laddr is not nil, it is
// used as the local address for the connection.
func DialTCP(net string, laddr, raddr *TCPAddr) (*TCPConn, error) {
	return dialTCP(net, laddr, raddr, noDeadline)
}

func dialTCP(net string, laddr, raddr *TCPAddr, deadline time.Time) (*TCPConn, error) {
	if !deadline.IsZero() {
		panic("net.dialTCP: deadline not implemented on Plan 9")
	}
	switch net {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, &OpError{"dial", net, raddr, UnknownNetworkError(net)}
	}
	if raddr == nil {
		return nil, &OpError{"dial", net, nil, errMissingAddress}
	}
	fd, err := dialPlan9(net, laddr, raddr)
	if err != nil {
		return nil, err
	}
	return newTCPConn(fd), nil
}

// TCPListener is a TCP network listener.  Clients should typically
// use variables of type Listener instead of assuming TCP.
type TCPListener struct {
	fd *netFD
}

// AcceptTCP accepts the next incoming call and returns the new
// connection and the remote address.
func (l *TCPListener) AcceptTCP() (*TCPConn, error) {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return nil, syscall.EINVAL
	}
	fd, err := l.fd.acceptPlan9()
	if err != nil {
		return nil, err
	}
	return newTCPConn(fd), nil
}

// Accept implements the Accept method in the Listener interface; it
// waits for the next call and returns a generic Conn.
func (l *TCPListener) Accept() (Conn, error) {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return nil, syscall.EINVAL
	}
	c, err := l.AcceptTCP()
	if err != nil {
		return nil, err
	}
	return c, nil
}

// Close stops listening on the TCP address.
// Already Accepted connections are not closed.
func (l *TCPListener) Close() error {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return syscall.EINVAL
	}
	if _, err := l.fd.ctl.WriteString("hangup"); err != nil {
		l.fd.ctl.Close()
		return &OpError{"close", l.fd.ctl.Name(), l.fd.laddr, err}
	}
	return l.fd.ctl.Close()
}

// Addr returns the listener's network address, a *TCPAddr.
func (l *TCPListener) Addr() Addr { return l.fd.laddr }

// SetDeadline sets the deadline associated with the listener.
// A zero time value disables the deadline.
func (l *TCPListener) SetDeadline(t time.Time) error {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return syscall.EINVAL
	}
	return setDeadline(l.fd, t)
}

// File returns a copy of the underlying os.File, set to blocking
// mode.  It is the caller's responsibility to close f when finished.
// Closing l does not affect f, and closing f does not affect l.
//
// The returned os.File's file descriptor is different from the
// connection's.  Attempting to change properties of the original
// using this duplicate may or may not have the desired effect.
func (l *TCPListener) File() (f *os.File, err error) { return l.dup() }

// ListenTCP announces on the TCP address laddr and returns a TCP
// listener.  Net must be "tcp", "tcp4", or "tcp6".  If laddr has a
// port of 0, ListenTCP will choose an available port.  The caller can
// use the Addr method of TCPListener to retrieve the chosen address.
func ListenTCP(net string, laddr *TCPAddr) (*TCPListener, error) {
	switch net {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, &OpError{"listen", net, laddr, UnknownNetworkError(net)}
	}
	if laddr == nil {
		laddr = &TCPAddr{}
	}
	fd, err := listenPlan9(net, laddr)
	if err != nil {
		return nil, err
	}
	return &TCPListener{fd}, nil
}
<commit_msg>net: add missing SetKeepAlivePeriod for Plan 9<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package net

import (
	"io"
	"os"
	"syscall"
	"time"
)

// TCPConn is an implementation of the Conn interface for TCP network
// connections.
type TCPConn struct {
	conn
}

func newTCPConn(fd *netFD) *TCPConn {
	return &TCPConn{conn{fd}}
}

// ReadFrom implements the io.ReaderFrom ReadFrom method.
func (c *TCPConn) ReadFrom(r io.Reader) (int64, error) {
	return genericReadFrom(c, r)
}

// CloseRead shuts down the reading side of the TCP connection.
// Most callers should just use Close.
func (c *TCPConn) CloseRead() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.CloseRead()
}

// CloseWrite shuts down the writing side of the TCP connection.
// Most callers should just use Close.
func (c *TCPConn) CloseWrite() error {
	if !c.ok() {
		return syscall.EINVAL
	}
	return c.fd.CloseWrite()
}

// SetLinger sets the behavior of Close() on a connection which still
// has data waiting to be sent or to be acknowledged.
//
// If sec < 0 (the default), Close returns immediately and the
// operating system finishes sending the data in the background.
//
// If sec == 0, Close returns immediately and the operating system
// discards any unsent or unacknowledged data.
//
// If sec > 0, Close blocks for at most sec seconds waiting for data
// to be sent and acknowledged.
func (c *TCPConn) SetLinger(sec int) error {
	return syscall.EPLAN9
}

// SetKeepAlive sets whether the operating system should send
// keepalive messages on the connection.
func (c *TCPConn) SetKeepAlive(keepalive bool) error {
	return syscall.EPLAN9
}

// SetKeepAlivePeriod sets period between keep alives.
func (c *TCPConn) SetKeepAlivePeriod(d time.Duration) error {
	return syscall.EPLAN9
}

// SetNoDelay controls whether the operating system should delay
// packet transmission in hopes of sending fewer packets (Nagle's
// algorithm).  The default is true (no delay), meaning that data is
// sent as soon as possible after a Write.
func (c *TCPConn) SetNoDelay(noDelay bool) error {
	return syscall.EPLAN9
}

// DialTCP connects to the remote address raddr on the network net,
// which must be "tcp", "tcp4", or "tcp6".  If laddr is not nil, it is
// used as the local address for the connection.
func DialTCP(net string, laddr, raddr *TCPAddr) (*TCPConn, error) {
	return dialTCP(net, laddr, raddr, noDeadline)
}

func dialTCP(net string, laddr, raddr *TCPAddr, deadline time.Time) (*TCPConn, error) {
	if !deadline.IsZero() {
		panic("net.dialTCP: deadline not implemented on Plan 9")
	}
	switch net {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, &OpError{"dial", net, raddr, UnknownNetworkError(net)}
	}
	if raddr == nil {
		return nil, &OpError{"dial", net, nil, errMissingAddress}
	}
	fd, err := dialPlan9(net, laddr, raddr)
	if err != nil {
		return nil, err
	}
	return newTCPConn(fd), nil
}

// TCPListener is a TCP network listener.  Clients should typically
// use variables of type Listener instead of assuming TCP.
type TCPListener struct {
	fd *netFD
}

// AcceptTCP accepts the next incoming call and returns the new
// connection and the remote address.
func (l *TCPListener) AcceptTCP() (*TCPConn, error) {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return nil, syscall.EINVAL
	}
	fd, err := l.fd.acceptPlan9()
	if err != nil {
		return nil, err
	}
	return newTCPConn(fd), nil
}

// Accept implements the Accept method in the Listener interface; it
// waits for the next call and returns a generic Conn.
func (l *TCPListener) Accept() (Conn, error) {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return nil, syscall.EINVAL
	}
	c, err := l.AcceptTCP()
	if err != nil {
		return nil, err
	}
	return c, nil
}

// Close stops listening on the TCP address.
// Already Accepted connections are not closed.
func (l *TCPListener) Close() error {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return syscall.EINVAL
	}
	if _, err := l.fd.ctl.WriteString("hangup"); err != nil {
		l.fd.ctl.Close()
		return &OpError{"close", l.fd.ctl.Name(), l.fd.laddr, err}
	}
	return l.fd.ctl.Close()
}

// Addr returns the listener's network address, a *TCPAddr.
func (l *TCPListener) Addr() Addr { return l.fd.laddr }

// SetDeadline sets the deadline associated with the listener.
// A zero time value disables the deadline.
func (l *TCPListener) SetDeadline(t time.Time) error {
	if l == nil || l.fd == nil || l.fd.ctl == nil {
		return syscall.EINVAL
	}
	return setDeadline(l.fd, t)
}

// File returns a copy of the underlying os.File, set to blocking
// mode.  It is the caller's responsibility to close f when finished.
// Closing l does not affect f, and closing f does not affect l.
//
// The returned os.File's file descriptor is different from the
// connection's.  Attempting to change properties of the original
// using this duplicate may or may not have the desired effect.
func (l *TCPListener) File() (f *os.File, err error) { return l.dup() }

// ListenTCP announces on the TCP address laddr and returns a TCP
// listener.  Net must be "tcp", "tcp4", or "tcp6".  If laddr has a
// port of 0, ListenTCP will choose an available port.  The caller can
// use the Addr method of TCPListener to retrieve the chosen address.
func ListenTCP(net string, laddr *TCPAddr) (*TCPListener, error) {
	switch net {
	case "tcp", "tcp4", "tcp6":
	default:
		return nil, &OpError{"listen", net, laddr, UnknownNetworkError(net)}
	}
	if laddr == nil {
		laddr = &TCPAddr{}
	}
	fd, err := listenPlan9(net, laddr)
	if err != nil {
		return nil, err
	}
	return &TCPListener{fd}, nil
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package exec

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"runtime"
	"strconv"
	"strings"
	"testing"
)

func helperCommand(s ...string) *Cmd {
	cs := []string{"-test.run=TestHelperProcess", "--"}
	cs = append(cs, s...)
	cmd := Command(os.Args[0], cs...)
	cmd.Env = append([]string{"GO_WANT_HELPER_PROCESS=1"}, os.Environ()...)
	return cmd
}

func TestEcho(t *testing.T) {
	bs, err := helperCommand("echo", "foo bar", "baz").Output()
	if err != nil {
		t.Errorf("echo: %v", err)
	}
	if g, e := string(bs), "foo bar baz\n"; g != e {
		t.Errorf("echo: want %q, got %q", e, g)
	}
}

func TestCatStdin(t *testing.T) {
	// Cat, testing stdin and stdout.
	input := "Input string\nLine 2"
	p := helperCommand("cat")
	p.Stdin = strings.NewReader(input)
	bs, err := p.Output()
	if err != nil {
		t.Errorf("cat: %v", err)
	}
	s := string(bs)
	if s != input {
		t.Errorf("cat: want %q, got %q", input, s)
	}
}

func TestCatGoodAndBadFile(t *testing.T) {
	// Testing combined output and error values.
	bs, err := helperCommand("cat", "/bogus/file.foo", "exec_test.go").CombinedOutput()
	if _, ok := err.(*ExitError); !ok {
		t.Errorf("expected *ExitError from cat combined; got %T: %v", err, err)
	}
	s := string(bs)
	sp := strings.SplitN(s, "\n", 2)
	if len(sp) != 2 {
		t.Fatalf("expected two lines from cat; got %q", s)
	}
	errLine, body := sp[0], sp[1]
	if !strings.HasPrefix(errLine, "Error: open /bogus/file.foo") {
		t.Errorf("expected stderr to complain about file; got %q", errLine)
	}
	if !strings.Contains(body, "func TestHelperProcess(t *testing.T)") {
		t.Errorf("expected test code; got %q (len %d)", body, len(body))
	}
}

func TestNoExistBinary(t *testing.T) {
	// Can't run a non-existent binary
	err := Command("/no-exist-binary").Run()
	if err == nil {
		t.Error("expected error from /no-exist-binary")
	}
}

func TestExitStatus(t *testing.T) {
	// Test that exit values are returned correctly
	err := helperCommand("exit", "42").Run()
	if werr, ok := err.(*ExitError); ok {
		if s, e := werr.Error(), "exit status 42"; s != e {
			t.Errorf("from exit 42 got exit %q, want %q", s, e)
		}
	} else {
		t.Fatalf("expected *ExitError from exit 42; got %T: %v", err, err)
	}
}

func TestPipes(t *testing.T) {
	check := func(what string, err error) {
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
	}
	// Cat, testing stdin and stdout.
	c := helperCommand("pipetest")
	stdin, err := c.StdinPipe()
	check("StdinPipe", err)
	stdout, err := c.StdoutPipe()
	check("StdoutPipe", err)
	stderr, err := c.StderrPipe()
	check("StderrPipe", err)

	outbr := bufio.NewReader(stdout)
	errbr := bufio.NewReader(stderr)
	line := func(what string, br *bufio.Reader) string {
		line, _, err := br.ReadLine()
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
		return string(line)
	}

	err = c.Start()
	check("Start", err)

	_, err = stdin.Write([]byte("O:I am output\n"))
	check("first stdin Write", err)
	if g, e := line("first output line", outbr), "O:I am output"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("E:I am error\n"))
	check("second stdin Write", err)
	if g, e := line("first error line", errbr), "E:I am error"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("O:I am output2\n"))
	check("third stdin Write 3", err)
	if g, e := line("second output line", outbr), "O:I am output2"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	stdin.Close()
	err = c.Wait()
	check("Wait", err)
}

func TestExtraFiles(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Logf("no operating system support; skipping")
		return
	}

	// Force network usage, to verify the epoll (or whatever) fd
	// doesn't leak to the child,
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	tf, err := ioutil.TempFile("", "")
	if err != nil {
		t.Fatalf("TempFile: %v", err)
	}
	defer os.Remove(tf.Name())
	defer tf.Close()

	const text = "Hello, fd 3!"
	_, err = tf.Write([]byte(text))
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	_, err = tf.Seek(0, os.SEEK_SET)
	if err != nil {
		t.Fatalf("Seek: %v", err)
	}

	c := helperCommand("read3")
	c.ExtraFiles = []*os.File{tf}
	bs, err := c.CombinedOutput()
	if err != nil {
		t.Fatalf("CombinedOutput: %v; output %q", err, bs)
	}
	if string(bs) != text {
		t.Errorf("got %q; want %q", string(bs), text)
	}
}

// TestHelperProcess isn't a real test. It's used as a helper process
// for TestParameterRun.
func TestHelperProcess(*testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	defer os.Exit(0)

	args := os.Args
	for len(args) > 0 {
		if args[0] == "--" {
			args = args[1:]
			break
		}
		args = args[1:]
	}
	if len(args) == 0 {
		fmt.Fprintf(os.Stderr, "No command\n")
		os.Exit(2)
	}

	cmd, args := args[0], args[1:]
	switch cmd {
	case "echo":
		iargs := []interface{}{}
		for _, s := range args {
			iargs = append(iargs, s)
		}
		fmt.Println(iargs...)
	case "cat":
		if len(args) == 0 {
			io.Copy(os.Stdout, os.Stdin)
			return
		}
		exit := 0
		for _, fn := range args {
			f, err := os.Open(fn)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				exit = 2
			} else {
				defer f.Close()
				io.Copy(os.Stdout, f)
			}
		}
		os.Exit(exit)
	case "pipetest":
		bufr := bufio.NewReader(os.Stdin)
		for {
			line, _, err := bufr.ReadLine()
			if err == io.EOF {
				break
			} else if err != nil {
				os.Exit(1)
			}
			if bytes.HasPrefix(line, []byte("O:")) {
				os.Stdout.Write(line)
				os.Stdout.Write([]byte{'\n'})
			} else if bytes.HasPrefix(line, []byte("E:")) {
				os.Stderr.Write(line)
				os.Stderr.Write([]byte{'\n'})
			} else {
				os.Exit(1)
			}
		}
	case "read3": // read fd 3
		fd3 := os.NewFile(3, "fd3")
		bs, err := ioutil.ReadAll(fd3)
		if err != nil {
			fmt.Printf("ReadAll from fd 3: %v", err)
			os.Exit(1)
		}
		// Now verify that there are no other open fds.
		var files []*os.File
		for wantfd := os.Stderr.Fd() + 2; wantfd <= 100; wantfd++ {
			f, err := os.Open(os.Args[0])
			if err != nil {
				fmt.Printf("error opening file with expected fd %d: %v", wantfd, err)
				os.Exit(1)
			}
			if got := f.Fd(); got != wantfd {
				fmt.Printf("leaked parent file. fd = %d; want %d", got, wantfd)
				os.Exit(1)
			}
			files = append(files, f)
		}
		for _, f := range files {
			f.Close()
		}
		os.Stderr.Write(bs)
	case "exit":
		n, _ := strconv.Atoi(args[0])
		os.Exit(n)
	default:
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
		os.Exit(2)
	}
}
<commit_msg>os/exec: dump lsof on failure<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package exec

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"runtime"
	"strconv"
	"strings"
	"testing"
)

func helperCommand(s ...string) *Cmd {
	cs := []string{"-test.run=TestHelperProcess", "--"}
	cs = append(cs, s...)
	cmd := Command(os.Args[0], cs...)
	cmd.Env = append([]string{"GO_WANT_HELPER_PROCESS=1"}, os.Environ()...)
	return cmd
}

func TestEcho(t *testing.T) {
	bs, err := helperCommand("echo", "foo bar", "baz").Output()
	if err != nil {
		t.Errorf("echo: %v", err)
	}
	if g, e := string(bs), "foo bar baz\n"; g != e {
		t.Errorf("echo: want %q, got %q", e, g)
	}
}

func TestCatStdin(t *testing.T) {
	// Cat, testing stdin and stdout.
	input := "Input string\nLine 2"
	p := helperCommand("cat")
	p.Stdin = strings.NewReader(input)
	bs, err := p.Output()
	if err != nil {
		t.Errorf("cat: %v", err)
	}
	s := string(bs)
	if s != input {
		t.Errorf("cat: want %q, got %q", input, s)
	}
}

func TestCatGoodAndBadFile(t *testing.T) {
	// Testing combined output and error values.
	bs, err := helperCommand("cat", "/bogus/file.foo", "exec_test.go").CombinedOutput()
	if _, ok := err.(*ExitError); !ok {
		t.Errorf("expected *ExitError from cat combined; got %T: %v", err, err)
	}
	s := string(bs)
	sp := strings.SplitN(s, "\n", 2)
	if len(sp) != 2 {
		t.Fatalf("expected two lines from cat; got %q", s)
	}
	errLine, body := sp[0], sp[1]
	if !strings.HasPrefix(errLine, "Error: open /bogus/file.foo") {
		t.Errorf("expected stderr to complain about file; got %q", errLine)
	}
	if !strings.Contains(body, "func TestHelperProcess(t *testing.T)") {
		t.Errorf("expected test code; got %q (len %d)", body, len(body))
	}
}

func TestNoExistBinary(t *testing.T) {
	// Can't run a non-existent binary
	err := Command("/no-exist-binary").Run()
	if err == nil {
		t.Error("expected error from /no-exist-binary")
	}
}

func TestExitStatus(t *testing.T) {
	// Test that exit values are returned correctly
	err := helperCommand("exit", "42").Run()
	if werr, ok := err.(*ExitError); ok {
		if s, e := werr.Error(), "exit status 42"; s != e {
			t.Errorf("from exit 42 got exit %q, want %q", s, e)
		}
	} else {
		t.Fatalf("expected *ExitError from exit 42; got %T: %v", err, err)
	}
}

func TestPipes(t *testing.T) {
	check := func(what string, err error) {
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
	}
	// Cat, testing stdin and stdout.
	c := helperCommand("pipetest")
	stdin, err := c.StdinPipe()
	check("StdinPipe", err)
	stdout, err := c.StdoutPipe()
	check("StdoutPipe", err)
	stderr, err := c.StderrPipe()
	check("StderrPipe", err)

	outbr := bufio.NewReader(stdout)
	errbr := bufio.NewReader(stderr)
	line := func(what string, br *bufio.Reader) string {
		line, _, err := br.ReadLine()
		if err != nil {
			t.Fatalf("%s: %v", what, err)
		}
		return string(line)
	}

	err = c.Start()
	check("Start", err)

	_, err = stdin.Write([]byte("O:I am output\n"))
	check("first stdin Write", err)
	if g, e := line("first output line", outbr), "O:I am output"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("E:I am error\n"))
	check("second stdin Write", err)
	if g, e := line("first error line", errbr), "E:I am error"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	_, err = stdin.Write([]byte("O:I am output2\n"))
	check("third stdin Write 3", err)
	if g, e := line("second output line", outbr), "O:I am output2"; g != e {
		t.Errorf("got %q, want %q", g, e)
	}

	stdin.Close()
	err = c.Wait()
	check("Wait", err)
}

func TestExtraFiles(t *testing.T) {
	if runtime.GOOS == "windows" {
		t.Logf("no operating system support; skipping")
		return
	}

	// Force network usage, to verify the epoll (or whatever) fd
	// doesn't leak to the child,
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	defer ln.Close()

	tf, err := ioutil.TempFile("", "")
	if err != nil {
		t.Fatalf("TempFile: %v", err)
	}
	defer os.Remove(tf.Name())
	defer tf.Close()

	const text = "Hello, fd 3!"
	_, err = tf.Write([]byte(text))
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	_, err = tf.Seek(0, os.SEEK_SET)
	if err != nil {
		t.Fatalf("Seek: %v", err)
	}

	c := helperCommand("read3")
	c.ExtraFiles = []*os.File{tf}
	bs, err := c.CombinedOutput()
	if err != nil {
		t.Fatalf("CombinedOutput: %v; output %q", err, bs)
	}
	if string(bs) != text {
		t.Errorf("got %q; want %q", string(bs), text)
	}
}

// TestHelperProcess isn't a real test. It's used as a helper process
// for TestParameterRun.
func TestHelperProcess(*testing.T) {
	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
		return
	}
	defer os.Exit(0)

	args := os.Args
	for len(args) > 0 {
		if args[0] == "--" {
			args = args[1:]
			break
		}
		args = args[1:]
	}
	if len(args) == 0 {
		fmt.Fprintf(os.Stderr, "No command\n")
		os.Exit(2)
	}

	cmd, args := args[0], args[1:]
	switch cmd {
	case "echo":
		iargs := []interface{}{}
		for _, s := range args {
			iargs = append(iargs, s)
		}
		fmt.Println(iargs...)
	case "cat":
		if len(args) == 0 {
			io.Copy(os.Stdout, os.Stdin)
			return
		}
		exit := 0
		for _, fn := range args {
			f, err := os.Open(fn)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				exit = 2
			} else {
				defer f.Close()
				io.Copy(os.Stdout, f)
			}
		}
		os.Exit(exit)
	case "pipetest":
		bufr := bufio.NewReader(os.Stdin)
		for {
			line, _, err := bufr.ReadLine()
			if err == io.EOF {
				break
			} else if err != nil {
				os.Exit(1)
			}
			if bytes.HasPrefix(line, []byte("O:")) {
				os.Stdout.Write(line)
				os.Stdout.Write([]byte{'\n'})
			} else if bytes.HasPrefix(line, []byte("E:")) {
				os.Stderr.Write(line)
				os.Stderr.Write([]byte{'\n'})
			} else {
				os.Exit(1)
			}
		}
	case "read3": // read fd 3
		fd3 := os.NewFile(3, "fd3")
		bs, err := ioutil.ReadAll(fd3)
		if err != nil {
			fmt.Printf("ReadAll from fd 3: %v", err)
			os.Exit(1)
		}
		// Now verify that there are no other open fds.
		var files []*os.File
		for wantfd := os.Stderr.Fd() + 2; wantfd <= 100; wantfd++ {
			f, err := os.Open(os.Args[0])
			if err != nil {
				fmt.Printf("error opening file with expected fd %d: %v", wantfd, err)
				fmt.Println(Command("lsof", "-p", fmt.Sprint(os.Getpid())).CombinedOutput())
				os.Exit(1)
			}
			if got := f.Fd(); got != wantfd {
				fmt.Printf("leaked parent file. fd = %d; want %d", got, wantfd)
				os.Exit(1)
			}
			files = append(files, f)
		}
		for _, f := range files {
			f.Close()
		}
		os.Stderr.Write(bs)
	case "exit":
		n, _ := strconv.Atoi(args[0])
		os.Exit(n)
	default:
		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
		os.Exit(2)
	}
}
<|endoftext|>


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

package set_test

import (
	"github.com/juju/names"
	"github.com/juju/testing"
	gc "launchpad.net/gocheck"

	"github.com/juju/utils/set"
)

type tagSetSuite struct {
	testing.IsolationSuite
}

var _ = gc.Suite(tagSetSuite{})

func (tagSetSuite) TestEmpty(c *gc.C) {
	t := set.NewTags()
	c.Assert(t.Size(), gc.Equals, 0)
}

func (tagSetSuite) TestInitialValues(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-rabbitmq-server-0")

	t := set.NewTags(foo, bar)
	c.Assert(t.Size(), gc.Equals, 2)
}

func (tagSetSuite) TestInitialStringValues(c *gc.C) {
	t, err := set.NewTagsFromStrings("unit-wordpress-0", "unit-rabbitmq-server-0")
	c.Assert(err, gc.IsNil)
	c.Assert(t.Size(), gc.Equals, 2)
}

func (tagSetSuite) TestSize(c *gc.C) {
	// Empty sets are empty.
	s := set.NewTags()
	c.Assert(s.Size(), gc.Equals, 0)

	// Size returns number of unique values.
	s, err := set.NewTagsFromStrings(
		"unit-wordpress-0",
		"unit-rabbitmq-server-0",
		"unit-wordpress-0",
	)
	c.Assert(err, gc.IsNil)
	c.Assert(s.Size(), gc.Equals, 2)
}

func (tagSetSuite) TestIsEmpty(c *gc.C) {
	// Empty sets are empty.
	s := set.NewTags()
	c.Assert(s.IsEmpty(), gc.Equals, true)

	// Non-empty sets are not empty.
	tag, _ := names.ParseTag("unit-wordpress-0")
	s = set.NewTags(tag)
	c.Assert(s.IsEmpty(), gc.Equals, false)

	// Newly empty sets work too.
	s.Remove(tag)
	c.Assert(s.IsEmpty(), gc.Equals, true)
}

func (tagSetSuite) TestAdd(c *gc.C) {
	t := set.NewTags()
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-rabbitmq-server-0")

	t.Add(foo)
	t.Add(bar)

	bar, _ = names.ParseTag("unit-wordpress-0")
	t.Add(bar)

	c.Assert(t.Size(), gc.Equals, 2)
}

func (tagSetSuite) TestRemove(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-rabbitmq-server-0")

	t := set.NewTags(foo, bar)
	t.Remove(foo)

	c.Assert(t.Contains(foo), gc.Equals, false)
	c.Assert(t.Contains(bar), gc.Equals, true)
}

func (tagSetSuite) TestContains(c *gc.C) {
	t, err := set.NewTagsFromStrings("unit-wordpress-0", "unit-rabbitmq-server-0")
	c.Assert(err, gc.IsNil)

	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-rabbitmq-server-0")
	baz, _ := names.ParseTag("unit-mongodb-0")

	c.Assert(t.Contains(foo), gc.Equals, true)
	c.Assert(t.Contains(bar), gc.Equals, true)
	c.Assert(t.Contains(baz), gc.Equals, false)
}

func (tagSetSuite) TestSortedValues(c *gc.C) {
	m1, _ := names.ParseTag("machine-0")
	z1, _ := names.ParseTag("unit-z-server-0")
	z2, _ := names.ParseTag("unit-z-server-1")
	a1, _ := names.ParseTag("unit-a-server-0")

	t := set.NewTags(z2, a1, z1, m1)
	values := t.SortedValues()

	c.Assert(values[0], gc.Equals, m1)
	c.Assert(values[1], gc.Equals, a1)
	c.Assert(values[2], gc.Equals, z1)
	c.Assert(values[3], gc.Equals, z2)
}

func (tagSetSuite) TestRemoveNonExistent(c *gc.C) {
	t := set.NewTags()
	foo, _ := names.ParseTag("unit-wordpress-0")
	t.Remove(foo)
	c.Assert(t.Size(), gc.Equals, 0)
}

func (tagSetSuite) TestUnion(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-mongodb-0")
	baz, _ := names.ParseTag("unit-rabbitmq-server-0")
	bang, _ := names.ParseTag("unit-mysql-server-0")

	t1 := set.NewTags(foo, bar)
	t2 := set.NewTags(foo, baz, bang)
	union1 := t1.Union(t2)
	union2 := t2.Union(t1)

	c.Assert(union1.Size(), gc.Equals, 4)
	c.Assert(union2.Size(), gc.Equals, 4)

	c.Assert(union1.Contains(foo), gc.Equals, true)
	c.Assert(union1.Contains(bar), gc.Equals, true)
	c.Assert(union1.Contains(baz), gc.Equals, true)
	c.Assert(union1.Contains(bang), gc.Equals, true)
}

func (tagSetSuite) TestIntersection(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-mongodb-0")
	baz, _ := names.ParseTag("unit-rabbitmq-server-0")
	bang, _ := names.ParseTag("unit-mysql-server-0")

	t1 := set.NewTags(foo, bar)
	t2 := set.NewTags(foo, baz, bang)

	int1 := t1.Intersection(t2)
	int2 := t2.Intersection(t1)

	c.Assert(int1.Size(), gc.Equals, 1)
	c.Assert(int2.Size(), gc.Equals, 1)

	c.Assert(int1.Contains(foo), gc.Equals, true)
	c.Assert(int2.Contains(foo), gc.Equals, true)
}

func (tagSetSuite) TestDifference(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-mongodb-0")
	baz, _ := names.ParseTag("unit-rabbitmq-server-0")
	bang, _ := names.ParseTag("unit-mysql-server-0")

	t1 := set.NewTags(foo, bar)
	t2 := set.NewTags(foo, baz, bang)

	diff1 := t1.Difference(t2)
	diff2 := t2.Difference(t1)

	c.Assert(diff1.Contains(bar), gc.Equals, true)
	c.Assert(diff2.Contains(baz), gc.Equals, true)
	c.Assert(diff2.Contains(bang), gc.Equals, true)

	c.Assert(diff1.Contains(foo), gc.Equals, false)
	c.Assert(diff2.Contains(foo), gc.Equals, false)
}

func (tagSetSuite) TestUninitialized(c *gc.C) {
	var uninitialized set.Tags

	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-mongodb-0")

	c.Assert(uninitialized.Size(), gc.Equals, 0)
	c.Assert(uninitialized.IsEmpty(), gc.Equals, true)
	// You can get values and sorted values from an unitialized set.
	c.Assert(uninitialized.Values(), gc.DeepEquals, []names.Tag{})
	// All contains checks are false
	c.Assert(uninitialized.Contains(foo), gc.Equals, false)
	// Remove works on an uninitialized Strings
	uninitialized.Remove(foo)

	var other set.Tags
	// Union returns a new set that is empty but initialized.
	c.Assert(uninitialized.Union(other), gc.DeepEquals, set.NewTags())
	c.Assert(uninitialized.Intersection(other), gc.DeepEquals, set.NewTags())
	c.Assert(uninitialized.Difference(other), gc.DeepEquals, set.NewTags())

	other = set.NewTags(foo, bar)
	c.Assert(uninitialized.Union(other), gc.DeepEquals, other)
	c.Assert(uninitialized.Intersection(other), gc.DeepEquals, set.NewTags())
	c.Assert(uninitialized.Difference(other), gc.DeepEquals, set.NewTags())
	c.Assert(other.Union(uninitialized), gc.DeepEquals, other)
	c.Assert(other.Intersection(uninitialized), gc.DeepEquals, set.NewTags())
	c.Assert(other.Difference(uninitialized), gc.DeepEquals, other)

	// Once something is added, the set becomes initialized.
	uninitialized.Add(foo)
	c.Assert(uninitialized.Contains(foo), gc.Equals, true)
}
<commit_msg>set/tags: update tests per review comments<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the LGPLv3, see LICENCE file for details.

package set_test

import (
	"github.com/juju/names"
	"github.com/juju/testing"
	gc "launchpad.net/gocheck"

	"github.com/juju/utils/set"
)

type tagSetSuite struct {
	testing.IsolationSuite
}

var _ = gc.Suite(tagSetSuite{})

func (tagSetSuite) TestEmpty(c *gc.C) {
	t := set.NewTags()
	c.Assert(t.Size(), gc.Equals, 0)
}

func (tagSetSuite) TestInitialValues(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-rabbitmq-server-0")

	t := set.NewTags(foo, bar)
	c.Assert(t.Size(), gc.Equals, 2)
}

func (tagSetSuite) TestInitialStringValues(c *gc.C) {
	t, err := set.NewTagsFromStrings("unit-wordpress-0", "unit-rabbitmq-server-0")
	c.Assert(err, gc.IsNil)
	c.Assert(t.Size(), gc.Equals, 2)
}

func (tagSetSuite) TestSize(c *gc.C) {
	// Empty sets are empty.
	s := set.NewTags()
	c.Assert(s.Size(), gc.Equals, 0)

	// Size returns number of unique values.
	s, err := set.NewTagsFromStrings(
		"unit-wordpress-0",
		"unit-rabbitmq-server-0",
		"unit-wordpress-0",
	)
	c.Assert(err, gc.IsNil)
	c.Assert(s.Size(), gc.Equals, 2)
}

func (tagSetSuite) TestIsEmpty(c *gc.C) {
	// Empty sets are empty.
	s := set.NewTags()
	c.Assert(s.IsEmpty(), gc.Equals, true)

	// Non-empty sets are not empty.
	tag, _ := names.ParseTag("unit-wordpress-0")
	s = set.NewTags(tag)
	c.Assert(s.IsEmpty(), gc.Equals, false)

	// Newly empty sets work too.
	s.Remove(tag)
	c.Assert(s.IsEmpty(), gc.Equals, true)
}

func (tagSetSuite) TestAdd(c *gc.C) {
	t := set.NewTags()
	foo, _ := names.ParseTag("unit-wordpress-0")
	t.Add(foo)
	c.Assert(t.Size(), gc.Equals, 1)
	c.Assert(t.Contains(foo), gc.Equals, true)
}

func (tagSetSuite) TestAddDuplicate(c *gc.C) {
	t := set.NewTags()
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-rabbitmq-server-0")

	t.Add(foo)
	t.Add(bar)

	bar, _ = names.ParseTag("unit-wordpress-0")
	t.Add(bar)

	c.Assert(t.Size(), gc.Equals, 2)
}

func (tagSetSuite) TestRemove(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-rabbitmq-server-0")

	t := set.NewTags(foo, bar)
	t.Remove(foo)

	c.Assert(t.Contains(foo), gc.Equals, false)
	c.Assert(t.Contains(bar), gc.Equals, true)
}

func (tagSetSuite) TestContains(c *gc.C) {
	t, err := set.NewTagsFromStrings("unit-wordpress-0", "unit-rabbitmq-server-0")
	c.Assert(err, gc.IsNil)

	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-rabbitmq-server-0")
	baz, _ := names.ParseTag("unit-mongodb-0")

	c.Assert(t.Contains(foo), gc.Equals, true)
	c.Assert(t.Contains(bar), gc.Equals, true)
	c.Assert(t.Contains(baz), gc.Equals, false)
}

func (tagSetSuite) TestSortedValues(c *gc.C) {
	m1, _ := names.ParseTag("machine-0")
	z1, _ := names.ParseTag("unit-z-server-0")
	z2, _ := names.ParseTag("unit-z-server-1")
	a1, _ := names.ParseTag("unit-a-server-0")

	t := set.NewTags(z2, a1, z1, m1)
	values := t.SortedValues()

	c.Assert(values, gc.DeepEquals, []names.Tag{m1, a1, z1, z2})
}

func (tagSetSuite) TestRemoveNonExistent(c *gc.C) {
	t := set.NewTags()
	foo, _ := names.ParseTag("unit-wordpress-0")
	t.Remove(foo)
	c.Assert(t.Size(), gc.Equals, 0)
}

func (tagSetSuite) TestUnion(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-mongodb-0")
	baz, _ := names.ParseTag("unit-rabbitmq-server-0")
	bang, _ := names.ParseTag("unit-mysql-server-0")

	t1 := set.NewTags(foo, bar)
	t2 := set.NewTags(foo, baz, bang)
	union1 := t1.Union(t2)
	union2 := t2.Union(t1)

	c.Assert(union1.Size(), gc.Equals, 4)
	c.Assert(union2.Size(), gc.Equals, 4)

	c.Assert(union1, gc.DeepEquals, union2)
	c.Assert(union1, gc.DeepEquals, set.NewTags(foo, bar, baz, bang))
}

func (tagSetSuite) TestIntersection(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-mongodb-0")
	baz, _ := names.ParseTag("unit-rabbitmq-server-0")
	bang, _ := names.ParseTag("unit-mysql-server-0")

	t1 := set.NewTags(foo, bar)
	t2 := set.NewTags(foo, baz, bang)

	int1 := t1.Intersection(t2)
	int2 := t2.Intersection(t1)

	c.Assert(int1.Size(), gc.Equals, 1)
	c.Assert(int2.Size(), gc.Equals, 1)

	c.Assert(int1, gc.DeepEquals, int2)
	c.Assert(int1, gc.DeepEquals, set.NewTags(foo))
}

func (tagSetSuite) TestDifference(c *gc.C) {
	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-mongodb-0")
	baz, _ := names.ParseTag("unit-rabbitmq-server-0")
	bang, _ := names.ParseTag("unit-mysql-server-0")

	t1 := set.NewTags(foo, bar)
	t2 := set.NewTags(foo, baz, bang)

	diff1 := t1.Difference(t2)
	diff2 := t2.Difference(t1)

	c.Assert(diff1, gc.DeepEquals, set.NewTags(bar))
	c.Assert(diff2, gc.DeepEquals, set.NewTags(baz, bang))
}

func (tagSetSuite) TestUninitialized(c *gc.C) {
	var uninitialized set.Tags

	foo, _ := names.ParseTag("unit-wordpress-0")
	bar, _ := names.ParseTag("unit-mongodb-0")

	c.Assert(uninitialized.Size(), gc.Equals, 0)
	c.Assert(uninitialized.IsEmpty(), gc.Equals, true)
	// You can get values and sorted values from an unitialized set.
	c.Assert(uninitialized.Values(), gc.DeepEquals, []names.Tag{})
	// All contains checks are false
	c.Assert(uninitialized.Contains(foo), gc.Equals, false)
	// Remove works on an uninitialized Strings
	uninitialized.Remove(foo)

	var other set.Tags
	// Union returns a new set that is empty but initialized.
	c.Assert(uninitialized.Union(other), gc.DeepEquals, set.NewTags())
	c.Assert(uninitialized.Intersection(other), gc.DeepEquals, set.NewTags())
	c.Assert(uninitialized.Difference(other), gc.DeepEquals, set.NewTags())

	other = set.NewTags(foo, bar)
	c.Assert(uninitialized.Union(other), gc.DeepEquals, other)
	c.Assert(uninitialized.Intersection(other), gc.DeepEquals, set.NewTags())
	c.Assert(uninitialized.Difference(other), gc.DeepEquals, set.NewTags())
	c.Assert(other.Union(uninitialized), gc.DeepEquals, other)
	c.Assert(other.Intersection(uninitialized), gc.DeepEquals, set.NewTags())
	c.Assert(other.Difference(uninitialized), gc.DeepEquals, other)

	// Once something is added, the set becomes initialized.
	uninitialized.Add(foo)
	c.Assert(uninitialized.Contains(foo), gc.Equals, true)
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

// Package tunny implements a simple pool for maintaining independant worker goroutines.
package tunny

import (
	"errors"
	"reflect"
	"sync"
	"sync/atomic"
	"time"
)

/*
TunnyWorker - The basic interface of a tunny worker.
*/
type TunnyWorker interface {

	// Called for each job, expects the result to be returned synchronously
	TunnyJob(interface{}) interface{}

	// Called after each job, this indicates whether the worker is ready for the next job.
	// The default implementation is to return true always. If false is returned then the
	// method is called every five milliseconds until either true is returned or the pool
	// is closed.
	TunnyReady() bool
}

/*
TunnyExtendedWorker - An optional interface that can be implemented if the worker needs
more control over its state.
*/
type TunnyExtendedWorker interface {

	// Called when the pool is opened, this will be called before any jobs are sent.
	TunnyInitialize()

	// Called when the pool is closed, this will be called after all jobs are completed.
	TunnyTerminate()
}

/*
TunnyInterruptable - An optional interface that can be implemented in order to allow the
worker to drop jobs when they are abandoned.
*/
type TunnyInterruptable interface {

	// Called when the current job has been abandoned by the client.
	TunnyInterrupt()
}

/*
Default and very basic implementation of a tunny worker. This worker holds a closure which
is assigned at construction, and this closure is called on each job.
*/
type tunnyDefaultWorker struct {
	job *func(interface{}) interface{}
}

func (worker *tunnyDefaultWorker) TunnyJob(data interface{}) interface{} {
	return (*worker.job)(data)
}

func (worker *tunnyDefaultWorker) TunnyReady() bool {
	return true
}

/*
WorkPool contains the structures and methods required to communicate with your pool, it must
be opened before sending work and closed when all jobs are completed.

You may open and close a pool as many times as you wish, calling close is a blocking call that
guarantees all goroutines are stopped.
*/
type WorkPool struct {
	workers     []*workerWrapper
	selects     []reflect.SelectCase
	statusMutex sync.RWMutex
	running     uint32
}

func (pool *WorkPool) isRunning() bool {
	return (atomic.LoadUint32(&pool.running) == 1)
}

func (pool *WorkPool) setRunning(running bool) {
	if running {
		atomic.SwapUint32(&pool.running, 1)
	} else {
		atomic.SwapUint32(&pool.running, 0)
	}
}

/*
Open all channels and launch the background goroutines managed by the pool.
*/
func (pool *WorkPool) Open() (*WorkPool, error) {
	pool.statusMutex.Lock()
	defer pool.statusMutex.Unlock()

	if !pool.isRunning() {

		pool.selects = make([]reflect.SelectCase, len(pool.workers))

		for i, workerWrapper := range pool.workers {
			workerWrapper.Open()

			pool.selects[i] = reflect.SelectCase{
				Dir:  reflect.SelectRecv,
				Chan: reflect.ValueOf(workerWrapper.readyChan),
			}
		}

		pool.setRunning(true)
		return pool, nil

	}
	return nil, errors.New("pool is already running")
}

/*
Close all channels and goroutines managed by the pool.
*/
func (pool *WorkPool) Close() error {
	pool.statusMutex.Lock()
	defer pool.statusMutex.Unlock()

	if pool.isRunning() {
		for _, workerWrapper := range pool.workers {
			workerWrapper.Close()
		}
		for _, workerWrapper := range pool.workers {
			workerWrapper.Join()
		}
		pool.setRunning(false)
		return nil
	}
	return errors.New("cannot close when the pool is not running")
}

/*
CreatePool - Creates a pool of workers, and takes a closure argument which is the action
to perform for each job.
*/
func CreatePool(numWorkers int, job func(interface{}) interface{}) *WorkPool {
	pool := WorkPool{running: 0}

	pool.workers = make([]*workerWrapper, numWorkers)
	for i := range pool.workers {
		newWorker := workerWrapper{
			worker: &(tunnyDefaultWorker{&job}),
		}
		pool.workers[i] = &newWorker
	}

	return &pool
}

/*
CreatePoolGeneric - Creates a pool of generic workers. When sending work to a pool of
generic workers you send a closure (func()) which is the job to perform.
*/
func CreatePoolGeneric(numWorkers int) *WorkPool {

	return CreatePool(numWorkers, func(jobCall interface{}) interface{} {
		if method, ok := jobCall.(func()); ok {
			method()
			return nil
		}
		return errors.New("generic worker not given a func()")
	})

}

/*
CreateCustomPool - Creates a pool for an array of custom workers. The custom workers
must implement TunnyWorker, and may also optionally implement TunnyExtendedWorker and
TunnyInterruptable.
*/
func CreateCustomPool(customWorkers []TunnyWorker) *WorkPool {
	pool := WorkPool{running: 0}

	pool.workers = make([]*workerWrapper, len(customWorkers))
	for i := range pool.workers {
		newWorker := workerWrapper{
			worker: customWorkers[i],
		}
		pool.workers[i] = &newWorker
	}

	return &pool
}

/*
SendWorkTimed - Send a job to a worker and return the result, this is a synchronous
call with a timeout.
*/
func (pool *WorkPool) SendWorkTimed(milliTimeout time.Duration, jobData interface{}) (interface{}, error) {
	pool.statusMutex.RLock()
	defer pool.statusMutex.RUnlock()

	if pool.isRunning() {
		before := time.Now()

		// Create new selectcase[] and add time out case
		selectCases := append(pool.selects[:], reflect.SelectCase{
			Dir:  reflect.SelectRecv,
			Chan: reflect.ValueOf(time.After(milliTimeout * time.Millisecond)),
		})

		// Wait for workers, or time out
		if chosen, _, ok := reflect.Select(selectCases); ok {

			if chosen < (len(selectCases) - 1) {
				pool.workers[chosen].jobChan <- jobData

				// Wait for response, or time out
				select {
				case data, open := <-pool.workers[chosen].outputChan:
					if !open {
						return nil, errors.New("worker was closed before reaching a result")
					}
					return data, nil
				case <-time.After((milliTimeout * time.Millisecond) - time.Since(before)):
					/* If we time out here we also need to ensure that the output is still
					 * collected and that the worker can move on. Therefore, we fork the
					 * waiting process into a new goroutine.
					 */
					go func() {
						pool.workers[chosen].Interrupt()
						<-pool.workers[chosen].outputChan
					}()
					return nil, errors.New("request timed out whilst waiting for job to complete")
				}
			} else {
				return nil, errors.New("request timed out whilst waiting for a worker")
			}
		} else {
			return nil, errors.New("failed to find a worker")
		}
	} else {
		return nil, errors.New("pool is not running! Call Open() before sending work")
	}
}

/*
SendWorkTimedAsync - Send a timed job to a worker without blocking, and optionally
send the result to a receiving closure. You may set the closure to nil if no
further actions are required.
*/
func (pool *WorkPool) SendWorkTimedAsync(
	milliTimeout time.Duration,
	jobData interface{},
	after func(interface{}, error),
) {
	go func() {
		result, err := pool.SendWorkTimed(milliTimeout, jobData)
		if after != nil {
			after(result, err)
		}
	}()
}

/*
SendWork - Send a job to a worker and return the result, this is a synchronous call.
*/
func (pool *WorkPool) SendWork(jobData interface{}) (interface{}, error) {
	pool.statusMutex.RLock()
	defer pool.statusMutex.RUnlock()

	if pool.isRunning() {
		if chosen, _, ok := reflect.Select(pool.selects); ok && chosen >= 0 {
			pool.workers[chosen].jobChan <- jobData
			result, open := <-pool.workers[chosen].outputChan

			if !open {
				return nil, errors.New("worker was closed before reaching a result")
			}
			return result, nil
		}
		return nil, errors.New("failed to find or wait for a worker")
	}
	return nil, errors.New("pool is not running! Call Open() before sending work")
}

/*
SendWorkAsync - Send a job to a worker without blocking, and optionally send the
result to a receiving closure. You may set the closure to nil if no further actions
are required.
*/
func (pool *WorkPool) SendWorkAsync(jobData interface{}, after func(interface{}, error)) {
	go func() {
		result, err := pool.SendWork(jobData)
		if after != nil {
			after(result, err)
		}
	}()
}
<commit_msg>Recycle and export errors<commit_after>/*
Copyright (c) 2014 Ashley Jeffs

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
*/

// Package tunny implements a simple pool for maintaining independant worker goroutines.
package tunny

import (
	"errors"
	"reflect"
	"sync"
	"sync/atomic"
	"time"
)

// Errors that are used throughout the Tunny API.
var (
	ErrPoolAlreadyRunning = errors.New("the pool is already running")
	ErrPoolNotRunning     = errors.New("the pool is not running")
	ErrJobNotFunc         = errors.New("generic worker not given a func()")
	ErrWorkerClosed       = errors.New("worker was closed")
	ErrJobTimedOut        = errors.New("job request timed out")
)

/*
TunnyWorker - The basic interface of a tunny worker.
*/
type TunnyWorker interface {

	// Called for each job, expects the result to be returned synchronously
	TunnyJob(interface{}) interface{}

	// Called after each job, this indicates whether the worker is ready for the next job.
	// The default implementation is to return true always. If false is returned then the
	// method is called every five milliseconds until either true is returned or the pool
	// is closed.
	TunnyReady() bool
}

/*
TunnyExtendedWorker - An optional interface that can be implemented if the worker needs
more control over its state.
*/
type TunnyExtendedWorker interface {

	// Called when the pool is opened, this will be called before any jobs are sent.
	TunnyInitialize()

	// Called when the pool is closed, this will be called after all jobs are completed.
	TunnyTerminate()
}

/*
TunnyInterruptable - An optional interface that can be implemented in order to allow the
worker to drop jobs when they are abandoned.
*/
type TunnyInterruptable interface {

	// Called when the current job has been abandoned by the client.
	TunnyInterrupt()
}

/*
Default and very basic implementation of a tunny worker. This worker holds a closure which
is assigned at construction, and this closure is called on each job.
*/
type tunnyDefaultWorker struct {
	job *func(interface{}) interface{}
}

func (worker *tunnyDefaultWorker) TunnyJob(data interface{}) interface{} {
	return (*worker.job)(data)
}

func (worker *tunnyDefaultWorker) TunnyReady() bool {
	return true
}

/*
WorkPool contains the structures and methods required to communicate with your pool, it must
be opened before sending work and closed when all jobs are completed.

You may open and close a pool as many times as you wish, calling close is a blocking call that
guarantees all goroutines are stopped.
*/
type WorkPool struct {
	workers     []*workerWrapper
	selects     []reflect.SelectCase
	statusMutex sync.RWMutex
	running     uint32
}

func (pool *WorkPool) isRunning() bool {
	return (atomic.LoadUint32(&pool.running) == 1)
}

func (pool *WorkPool) setRunning(running bool) {
	if running {
		atomic.SwapUint32(&pool.running, 1)
	} else {
		atomic.SwapUint32(&pool.running, 0)
	}
}

/*
Open all channels and launch the background goroutines managed by the pool.
*/
func (pool *WorkPool) Open() (*WorkPool, error) {
	pool.statusMutex.Lock()
	defer pool.statusMutex.Unlock()

	if !pool.isRunning() {

		pool.selects = make([]reflect.SelectCase, len(pool.workers))

		for i, workerWrapper := range pool.workers {
			workerWrapper.Open()

			pool.selects[i] = reflect.SelectCase{
				Dir:  reflect.SelectRecv,
				Chan: reflect.ValueOf(workerWrapper.readyChan),
			}
		}

		pool.setRunning(true)
		return pool, nil

	}
	return nil, ErrPoolAlreadyRunning
}

/*
Close all channels and goroutines managed by the pool.
*/
func (pool *WorkPool) Close() error {
	pool.statusMutex.Lock()
	defer pool.statusMutex.Unlock()

	if pool.isRunning() {
		for _, workerWrapper := range pool.workers {
			workerWrapper.Close()
		}
		for _, workerWrapper := range pool.workers {
			workerWrapper.Join()
		}
		pool.setRunning(false)
		return nil
	}
	return ErrPoolNotRunning
}

/*
CreatePool - Creates a pool of workers, and takes a closure argument which is the action
to perform for each job.
*/
func CreatePool(numWorkers int, job func(interface{}) interface{}) *WorkPool {
	pool := WorkPool{running: 0}

	pool.workers = make([]*workerWrapper, numWorkers)
	for i := range pool.workers {
		newWorker := workerWrapper{
			worker: &(tunnyDefaultWorker{&job}),
		}
		pool.workers[i] = &newWorker
	}

	return &pool
}

/*
CreatePoolGeneric - Creates a pool of generic workers. When sending work to a pool of
generic workers you send a closure (func()) which is the job to perform.
*/
func CreatePoolGeneric(numWorkers int) *WorkPool {

	return CreatePool(numWorkers, func(jobCall interface{}) interface{} {
		if method, ok := jobCall.(func()); ok {
			method()
			return nil
		}
		return ErrJobNotFunc
	})

}

/*
CreateCustomPool - Creates a pool for an array of custom workers. The custom workers
must implement TunnyWorker, and may also optionally implement TunnyExtendedWorker and
TunnyInterruptable.
*/
func CreateCustomPool(customWorkers []TunnyWorker) *WorkPool {
	pool := WorkPool{running: 0}

	pool.workers = make([]*workerWrapper, len(customWorkers))
	for i := range pool.workers {
		newWorker := workerWrapper{
			worker: customWorkers[i],
		}
		pool.workers[i] = &newWorker
	}

	return &pool
}

/*
SendWorkTimed - Send a job to a worker and return the result, this is a synchronous
call with a timeout.
*/
func (pool *WorkPool) SendWorkTimed(milliTimeout time.Duration, jobData interface{}) (interface{}, error) {
	pool.statusMutex.RLock()
	defer pool.statusMutex.RUnlock()

	if pool.isRunning() {
		before := time.Now()

		// Create new selectcase[] and add time out case
		selectCases := append(pool.selects[:], reflect.SelectCase{
			Dir:  reflect.SelectRecv,
			Chan: reflect.ValueOf(time.After(milliTimeout * time.Millisecond)),
		})

		// Wait for workers, or time out
		if chosen, _, ok := reflect.Select(selectCases); ok {

			// Check if the selected index is a worker, otherwise we timed out
			if chosen < (len(selectCases) - 1) {
				pool.workers[chosen].jobChan <- jobData

				// Wait for response, or time out
				select {
				case data, open := <-pool.workers[chosen].outputChan:
					if !open {
						return nil, ErrWorkerClosed
					}
					return data, nil
				case <-time.After((milliTimeout * time.Millisecond) - time.Since(before)):
					/* If we time out here we also need to ensure that the output is still
					 * collected and that the worker can move on. Therefore, we fork the
					 * waiting process into a new goroutine.
					 */
					go func() {
						pool.workers[chosen].Interrupt()
						<-pool.workers[chosen].outputChan
					}()
					return nil, ErrJobTimedOut
				}
			} else {
				return nil, ErrJobTimedOut
			}
		} else {
			// This means the chosen channel was closed
			return nil, ErrWorkerClosed
		}
	} else {
		return nil, ErrPoolNotRunning
	}
}

/*
SendWorkTimedAsync - Send a timed job to a worker without blocking, and optionally
send the result to a receiving closure. You may set the closure to nil if no
further actions are required.
*/
func (pool *WorkPool) SendWorkTimedAsync(
	milliTimeout time.Duration,
	jobData interface{},
	after func(interface{}, error),
) {
	go func() {
		result, err := pool.SendWorkTimed(milliTimeout, jobData)
		if after != nil {
			after(result, err)
		}
	}()
}

/*
SendWork - Send a job to a worker and return the result, this is a synchronous call.
*/
func (pool *WorkPool) SendWork(jobData interface{}) (interface{}, error) {
	pool.statusMutex.RLock()
	defer pool.statusMutex.RUnlock()

	if pool.isRunning() {
		if chosen, _, ok := reflect.Select(pool.selects); ok && chosen >= 0 {
			pool.workers[chosen].jobChan <- jobData
			result, open := <-pool.workers[chosen].outputChan

			if !open {
				return nil, ErrWorkerClosed
			}
			return result, nil
		}
		return nil, ErrWorkerClosed
	}
	return nil, ErrPoolNotRunning
}

/*
SendWorkAsync - Send a job to a worker without blocking, and optionally send the
result to a receiving closure. You may set the closure to nil if no further actions
are required.
*/
func (pool *WorkPool) SendWorkAsync(jobData interface{}, after func(interface{}, error)) {
	go func() {
		result, err := pool.SendWork(jobData)
		if after != nil {
			after(result, err)
		}
	}()
}
<|endoftext|>


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

package samples

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os/exec"
	"path"
	"sync"

	"github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

// A struct that implements common behavior needed by tests in the samples/
// directory where the file system is mounted by a subprocess. Use it as an
// embedded field in your test fixture, calling its SetUp method from your
// SetUp method after setting the MountType and MountFlags fields.
type SubprocessTest struct {
	// The type of the file system to mount. Must be recognized by mount_sample.
	MountType string

	// Additional flags to be passed to the mount_sample tool.
	MountFlags []string

	// A context object that can be used for long-running operations.
	Ctx context.Context

	// The directory at which the file system is mounted.
	Dir string

	// Anothing non-nil in this slice will be closed by TearDown. The test will
	// fail if closing fails.
	ToClose []io.Closer

	mountCmd *exec.Cmd
}

// Mount the file system and initialize the other exported fields of the
// struct. Panics on error.
//
// REQUIRES: t.FileSystem has been set.
func (t *SubprocessTest) SetUp(ti *ogletest.TestInfo) {
	err := t.initialize()
	if err != nil {
		panic(err)
	}
}

// Set by buildMountSample.
var mountSamplePath string
var mountSampleErr error
var mountSampleOnce sync.Once

// Build the mount_sample tool if it has not yet been built for this process.
// Return a path to the binary.
func buildMountSample() (toolPath string, err error) {
	// Build if we haven't yet.
	mountSampleOnce.Do(func() {
		// Create a temporary directory.
		tempDir, err := ioutil.TempDir("", "")
		if err != nil {
			mountSampleErr = fmt.Errorf("TempDir: %v", err)
			return
		}

		mountSamplePath = path.Join(tempDir, "mount_sample")

		// Build the command.
		cmd := exec.Command(
			"go",
			"build",
			"-o",
			mountSamplePath,
			"github.com/jacobsa/fuse/samples/mount_sample")

		output, err := cmd.CombinedOutput()
		if err != nil {
			mountSampleErr = fmt.Errorf(
				"mount_sample exited with %v, output:\n%s",
				err,
				string(output))

			return
		}
	})

	if mountSampleErr != nil {
		err = mountSampleErr
		return
	}

	toolPath = mountSamplePath
	return
}

// Invoke mount_sample, returning a running command.
func invokeMountSample(path string, args []string) (cmd *exec.Cmd, err error) {
	cmd = exec.Command(path, args...)
	if err = cmd.Start(); err != nil {
		err = fmt.Errorf("Start: %v", err)
		return
	}

	return
}

// Like SetUp, but doens't panic.
func (t *SubprocessTest) initialize() (err error) {
	// Initialize the context.
	t.Ctx = context.Background()

	// Set up a temporary directory.
	t.Dir, err = ioutil.TempDir("", "sample_test")
	if err != nil {
		err = fmt.Errorf("TempDir: %v", err)
		return
	}

	// Build the mount_sample tool.
	toolPath, err := buildMountSample()
	if err != nil {
		err = fmt.Errorf("buildMountSample: %v", err)
		return
	}

	// Invoke it.
	args := []string{"--type", t.MountType}
	args = append(args, t.MountFlags...)

	t.mountCmd, err = invokeMountSample(toolPath, args)
	if err != nil {
		err = fmt.Errorf("invokeMountSample: %v", err)
		return
	}

	// TODO(jacobsa): Probably need some sort of signalling (on stderr? write to
	// a flag-controlled file?) when WaitForReady has returned.

	return
}

// Unmount the file system and clean up. Panics on error.
func (t *SubprocessTest) TearDown() {
	err := t.destroy()
	if err != nil {
		panic(err)
	}
}

// Like TearDown, but doesn't panic.
func (t *SubprocessTest) destroy() (err error) {
	// Close what is necessary.
	for _, c := range t.ToClose {
		if c == nil {
			continue
		}

		ogletest.ExpectEq(nil, c.Close())
	}

	// If we didn't try to mount the file system, there's nothing further to do.
	if t.mountCmd == nil {
		return
	}

	// In the background, initiate an unmount.
	unmountErrChan := make(chan error)
	go func() {
		unmountErrChan <- unmount(t.Dir)
	}()

	// Make sure we wait for the unmount, even if we've already returned early in
	// error. Return its error if we haven't seen any other error.
	defer func() {
		unmountErr := <-unmountErrChan
		if unmountErr != nil {
			if err != nil {
				log.Println("unmount:", unmountErr)
				return
			}

			err = fmt.Errorf("unmount: %v", unmountErr)
		}
	}()

	// Wait for the subprocess.
	if err = t.mountCmd.Wait(); err != nil {
		err = fmt.Errorf("Cmd.Wait: %v", err)
		return
	}

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

package samples

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os/exec"
	"path"
	"sync"

	"github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
)

// A struct that implements common behavior needed by tests in the samples/
// directory where the file system is mounted by a subprocess. Use it as an
// embedded field in your test fixture, calling its SetUp method from your
// SetUp method after setting the MountType and MountFlags fields.
type SubprocessTest struct {
	// The type of the file system to mount. Must be recognized by mount_sample.
	MountType string

	// Additional flags to be passed to the mount_sample tool.
	MountFlags []string

	// A context object that can be used for long-running operations.
	Ctx context.Context

	// The directory at which the file system is mounted.
	Dir string

	// Anothing non-nil in this slice will be closed by TearDown. The test will
	// fail if closing fails.
	ToClose []io.Closer

	mountCmd    *exec.Cmd
	mountStdout bytes.Buffer
	mountStderr bytes.Buffer
}

// Mount the file system and initialize the other exported fields of the
// struct. Panics on error.
//
// REQUIRES: t.FileSystem has been set.
func (t *SubprocessTest) SetUp(ti *ogletest.TestInfo) {
	err := t.initialize()
	if err != nil {
		panic(err)
	}
}

// Set by buildMountSample.
var mountSamplePath string
var mountSampleErr error
var mountSampleOnce sync.Once

// Build the mount_sample tool if it has not yet been built for this process.
// Return a path to the binary.
func buildMountSample() (toolPath string, err error) {
	// Build if we haven't yet.
	mountSampleOnce.Do(func() {
		// Create a temporary directory.
		tempDir, err := ioutil.TempDir("", "")
		if err != nil {
			mountSampleErr = fmt.Errorf("TempDir: %v", err)
			return
		}

		mountSamplePath = path.Join(tempDir, "mount_sample")

		// Build the command.
		cmd := exec.Command(
			"go",
			"build",
			"-o",
			mountSamplePath,
			"github.com/jacobsa/fuse/samples/mount_sample")

		output, err := cmd.CombinedOutput()
		if err != nil {
			mountSampleErr = fmt.Errorf(
				"mount_sample exited with %v, output:\n%s",
				err,
				string(output))

			return
		}
	})

	if mountSampleErr != nil {
		err = mountSampleErr
		return
	}

	toolPath = mountSamplePath
	return
}

// Like SetUp, but doens't panic.
func (t *SubprocessTest) initialize() (err error) {
	// Initialize the context.
	t.Ctx = context.Background()

	// Set up a temporary directory.
	t.Dir, err = ioutil.TempDir("", "sample_test")
	if err != nil {
		err = fmt.Errorf("TempDir: %v", err)
		return
	}

	// Build the mount_sample tool.
	toolPath, err := buildMountSample()
	if err != nil {
		err = fmt.Errorf("buildMountSample: %v", err)
		return
	}

	// Set up a command.
	args := []string{
		toolPath,
		"--type",
		t.MountType,
	}

	args = append(args, t.MountFlags...)

	t.mountCmd = exec.Command(toolPath, args...)
	t.mountCmd.Stdout = &t.mountStdout
	t.mountCmd.Stderr = &t.mountStderr

	// Start it.
	if err = t.mountCmd.Start(); err != nil {
		err = fmt.Errorf("mountCmd.Start: %v", err)
		return
	}

	// TODO(jacobsa): Probably need some sort of signalling (on stderr? write to
	// a flag-controlled file?) when WaitForReady has returned.

	return
}

// Unmount the file system and clean up. Panics on error.
func (t *SubprocessTest) TearDown() {
	err := t.destroy()
	if err != nil {
		panic(err)
	}
}

// Like TearDown, but doesn't panic.
func (t *SubprocessTest) destroy() (err error) {
	// Close what is necessary.
	for _, c := range t.ToClose {
		if c == nil {
			continue
		}

		ogletest.ExpectEq(nil, c.Close())
	}

	// If we didn't try to mount the file system, there's nothing further to do.
	if t.mountCmd == nil {
		return
	}

	// In the background, initiate an unmount.
	unmountErrChan := make(chan error)
	go func() {
		unmountErrChan <- unmount(t.Dir)
	}()

	// Make sure we wait for the unmount, even if we've already returned early in
	// error. Return its error if we haven't seen any other error.
	defer func() {
		unmountErr := <-unmountErrChan
		if unmountErr != nil {
			if err != nil {
				log.Println("unmount:", unmountErr)
				return
			}

			err = fmt.Errorf("unmount: %v", unmountErr)
		}
	}()

	// Wait for the subprocess.
	if err = t.mountCmd.Wait(); err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			err = fmt.Errorf(
				"mount_sample exited with %v. Stderr:\n%s",
				exitErr,
				t.mountStderr.String())

			return
		}

		err = fmt.Errorf("mountCmd.Wait: %v", err)
		return
	}

	return
}
<|endoftext|>