<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/jingweno/travisarchive/db"
)

var (
	newBuildCrawlerInterval      = 30 * time.Second
	finishedBuildCrawlerInterval = 2 * time.Minute
)

type Crawler interface {
	Crawl()
}

func NewCrawler(travis *Travis, db *db.DB) []Crawler {
	return []Crawler{
		&NewBuildCrawler{travis, db, log.New(os.Stderr, "[NewBuildCrawler] ", log.LstdFlags)},
		&FinishedBuildCrawler{travis, db, log.New(os.Stderr, "[FinishedBuildCrawler] ", log.LstdFlags)},
	}
}

type NewBuildCrawler struct {
	Travis *Travis
	DB     *db.DB
	Logger *log.Logger
}

func (c *NewBuildCrawler) Crawl() {
	ch := time.Tick(newBuildCrawlerInterval)
	for _ = range ch {
		c.Logger.Println("crawling for new builds...")
		c.crawlNewBuilds()
	}
}

func (c *NewBuildCrawler) crawlNewBuilds() {
	repos, err := c.Travis.Repos()
	if err != nil {
		c.Logger.Println(err)
		return
	}

	newBuilds := []string{}
	for _, repo := range repos {
		updated, err := c.DB.Upsert("new_builds", db.Query{"lastbuildid": repo.LastBuildId}, repo)
		if err != nil {
			c.Logger.Println(err)
			continue
		}

		if updated {
			newBuilds = append(newBuilds, repo.Slug)
			err := stathat.CountOne(repo.LastBuildStartedAt)
			if err != nil {
				c.Logger.Println(err)
			}
		}
	}

	c.Logger.Printf("harvested %d builds with %d new builds: %s\n", len(repos), len(newBuilds), strings.Join(newBuilds, ", "))
}

type FinishedBuildCrawler struct {
	Travis *Travis
	DB     *db.DB
	Logger *log.Logger
}

func (c *FinishedBuildCrawler) Crawl() {
	ch := time.Tick(finishedBuildCrawlerInterval)
	for _ = range ch {
		c.Logger.Println("crawling for finsihed builds...")
		c.crawlFinishedBuilds()
	}
}

func (c *FinishedBuildCrawler) crawlFinishedBuilds() {
	colNames, finishedBuilds, skippedBuilds := c.doCrawlFinishedBuilds()
	c.Logger.Printf("fetched %d builds with %d finsihed and %d skipped. Finsihed builds: %s\n", len(finishedBuilds)+len(skippedBuilds), len(finishedBuilds), len(skippedBuilds), strings.Join(finishedBuilds, ", "))

	err := c.ensureColIndexes(colNames)
	if err != nil {
		c.Logger.Println(err)
	}
}

func (c *FinishedBuildCrawler) doCrawlFinishedBuilds() (colNames map[string]string, finishedBuilds []string, skippedBuilds []string) {
	colNames = make(map[string]string)

	var (
		repo *Repo
		//query db.Query
	)
	//query = Query{"lastbuildstartedat": Query{"$gte": oneMinuteAgo()}}
	iter := c.DB.C("new_builds").Find(nil).Sort("lastbuildstartedat").Iter()
	for iter.Next(&repo) {
		build, err := c.crawlFinsihedBuild(repo)
		if err != nil {
			// this build has real error
			// otherwise it's still in progress
			if build == nil {
				c.Logger.Println(err)
			}
			skippedBuilds = append(skippedBuilds, repo.Slug)
			continue
		}

		err = c.logBuild(build)
		if err != nil {
			c.Logger.Println(err)
			continue
		}

		colName, updated, err := c.upsertBuild(build)
		if err != nil {
			c.Logger.Println(err)
			continue
		}

		colNames[colName] = colName

		if updated {
			finishedBuilds = append(finishedBuilds, repo.Slug)
		}
	}

	if err := iter.Close(); err != nil {
		c.Logger.Println(err)
	}

	return
}

func (c *FinishedBuildCrawler) logBuild(build *Build) error {
	b, err := json.Marshal(build)
	if err != nil {
		return err
	}

	fmt.Printf("@[production.build] %s\n", b)

	return nil
}

func (c *FinishedBuildCrawler) crawlFinsihedBuild(repo *Repo) (build *Build, err error) {
	build, err = c.Travis.Build(repo.LastBuildId)
	if err != nil {
		return
	}

	isFinished := build.FinishedAt != nil && build.StartedAt != nil
	if !isFinished {
		err = fmt.Errorf("skipping build: %s - %d\n", repo.Slug, repo.LastBuildId)
		return
	}

	build.Repository = repo

	return
}

func (c *FinishedBuildCrawler) upsertBuild(build *Build) (colName string, updated bool, err error) {
	colName = buildColName(build.StartedAt)

	updated, err = c.DB.Upsert(colName, db.Query{"id": build.Id}, build)
	if err != nil {
		return
	}

	err = c.DB.C("new_builds").Remove(db.Query{"lastbuildid": build.Repository.LastBuildId})
	if err != nil {
		return
	}

	return
}

func (c *FinishedBuildCrawler) ensureColIndexes(colNames map[string]string) error {
	for _, colName := range colNames {
		c.Logger.Printf("ensuring index for collection %s\n", colName)
		err := c.DB.EnsureUniqueIndexKey(colName, "id")
		if err != nil {
			return err
		}
	}

	return nil
}

func buildColName(date *time.Time) string {
	buildDate := date.UTC().Format("2006_01_02")
	return fmt.Sprintf("builds_%s", buildDate)
}

func oneMinuteAgo() time.Time {
	now := time.Now()
	return now.Add(-1 * time.Minute).UTC()
}
<commit_msg>Remove log build since it doesn't work<commit_after>package main

import (
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"github.com/jingweno/travisarchive/db"
)

var (
	newBuildCrawlerInterval      = 30 * time.Second
	finishedBuildCrawlerInterval = 2 * time.Minute
)

type Crawler interface {
	Crawl()
}

func NewCrawler(travis *Travis, db *db.DB) []Crawler {
	return []Crawler{
		&NewBuildCrawler{travis, db, log.New(os.Stderr, "[NewBuildCrawler] ", log.LstdFlags)},
		&FinishedBuildCrawler{travis, db, log.New(os.Stderr, "[FinishedBuildCrawler] ", log.LstdFlags)},
	}
}

type NewBuildCrawler struct {
	Travis *Travis
	DB     *db.DB
	Logger *log.Logger
}

func (c *NewBuildCrawler) Crawl() {
	ch := time.Tick(newBuildCrawlerInterval)
	for _ = range ch {
		c.Logger.Println("crawling for new builds...")
		c.crawlNewBuilds()
	}
}

func (c *NewBuildCrawler) crawlNewBuilds() {
	repos, err := c.Travis.Repos()
	if err != nil {
		c.Logger.Println(err)
		return
	}

	newBuilds := []string{}
	for _, repo := range repos {
		updated, err := c.DB.Upsert("new_builds", db.Query{"lastbuildid": repo.LastBuildId}, repo)
		if err != nil {
			c.Logger.Println(err)
			continue
		}

		if updated {
			newBuilds = append(newBuilds, repo.Slug)
			err := stathat.CountOne(repo.LastBuildStartedAt)
			if err != nil {
				c.Logger.Println(err)
			}
		}
	}

	c.Logger.Printf("harvested %d builds with %d new builds: %s\n", len(repos), len(newBuilds), strings.Join(newBuilds, ", "))
}

type FinishedBuildCrawler struct {
	Travis *Travis
	DB     *db.DB
	Logger *log.Logger
}

func (c *FinishedBuildCrawler) Crawl() {
	ch := time.Tick(finishedBuildCrawlerInterval)
	for _ = range ch {
		c.Logger.Println("crawling for finsihed builds...")
		c.crawlFinishedBuilds()
	}
}

func (c *FinishedBuildCrawler) crawlFinishedBuilds() {
	colNames, finishedBuilds, skippedBuilds := c.doCrawlFinishedBuilds()
	c.Logger.Printf("fetched %d builds with %d finsihed and %d skipped. Finsihed builds: %s\n", len(finishedBuilds)+len(skippedBuilds), len(finishedBuilds), len(skippedBuilds), strings.Join(finishedBuilds, ", "))

	err := c.ensureColIndexes(colNames)
	if err != nil {
		c.Logger.Println(err)
	}
}

func (c *FinishedBuildCrawler) doCrawlFinishedBuilds() (colNames map[string]string, finishedBuilds []string, skippedBuilds []string) {
	colNames = make(map[string]string)

	var (
		repo *Repo
		//query db.Query
	)
	//query = Query{"lastbuildstartedat": Query{"$gte": oneMinuteAgo()}}
	iter := c.DB.C("new_builds").Find(nil).Sort("lastbuildstartedat").Iter()
	for iter.Next(&repo) {
		build, err := c.crawlFinsihedBuild(repo)
		if err != nil {
			// this build has real error
			// otherwise it's still in progress
			if build == nil {
				c.Logger.Println(err)
			}
			skippedBuilds = append(skippedBuilds, repo.Slug)
			continue
		}

		colName, updated, err := c.upsertBuild(build)
		if err != nil {
			c.Logger.Println(err)
			continue
		}

		colNames[colName] = colName

		if updated {
			finishedBuilds = append(finishedBuilds, repo.Slug)
		}
	}

	if err := iter.Close(); err != nil {
		c.Logger.Println(err)
	}

	return
}

func (c *FinishedBuildCrawler) crawlFinsihedBuild(repo *Repo) (build *Build, err error) {
	build, err = c.Travis.Build(repo.LastBuildId)
	if err != nil {
		return
	}

	isFinished := build.FinishedAt != nil && build.StartedAt != nil
	if !isFinished {
		err = fmt.Errorf("skipping build: %s - %d\n", repo.Slug, repo.LastBuildId)
		return
	}

	build.Repository = repo

	return
}

func (c *FinishedBuildCrawler) upsertBuild(build *Build) (colName string, updated bool, err error) {
	colName = buildColName(build.StartedAt)

	updated, err = c.DB.Upsert(colName, db.Query{"id": build.Id}, build)
	if err != nil {
		return
	}

	err = c.DB.C("new_builds").Remove(db.Query{"lastbuildid": build.Repository.LastBuildId})
	if err != nil {
		return
	}

	return
}

func (c *FinishedBuildCrawler) ensureColIndexes(colNames map[string]string) error {
	for _, colName := range colNames {
		c.Logger.Printf("ensuring index for collection %s\n", colName)
		err := c.DB.EnsureUniqueIndexKey(colName, "id")
		if err != nil {
			return err
		}
	}

	return nil
}

func buildColName(date *time.Time) string {
	buildDate := date.UTC().Format("2006_01_02")
	return fmt.Sprintf("builds_%s", buildDate)
}

func oneMinuteAgo() time.Time {
	now := time.Now()
	return now.Add(-1 * time.Minute).UTC()
}
<|endoftext|>


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

package main

import (
	"fmt"
	"math"

	"github.com/Jeffail/gabs"
)

type CurseForgeModFile struct {
	projectID  int
	fileID     int
	desc       string
	name       string
	clientOnly bool
}

func SelectCurseForgeModFile(pack *ModPack, mod string, url string, clientOnly bool) error {
	// Try to find the project ID using the mod name as a slug
	projectID, err := pack.db.findModBySlug(mod)
	if err != nil {
		return fmt.Errorf("unknown mod %s", mod)
	}

	// Look up the slug, name and description
	_, name, desc, err := pack.db.getProjectInfo(projectID)
	if err != nil {
		return fmt.Errorf("no name/description available for %s (%d): %+v", mod, projectID, err)
	}

	// Setup a mod file entry and then pull the latest file info
	modFile := CurseForgeModFile{projectID: projectID, desc: desc, name: name, clientOnly: clientOnly}
	fileId, err := modFile.getLatestFile(pack.minecraftVersion())
	if err != nil {
		return fmt.Errorf("failed to get latest file for %s (%d): %+v", mod, projectID, err)
	}

	// If we found a newer file, update entry and then the pack
	if fileId > modFile.fileID {
		modFile.fileID = fileId
		err = pack.selectMod(&modFile)
		if err != nil {
			return err
		}
	}

	return nil
}

func NewCurseForgeModFile(modJson *gabs.Container) *CurseForgeModFile {
	projectID, _ := intValue(modJson, "projectID")
	fileID, _ := intValue(modJson, "fileID")
	name, ok := modJson.Path("desc").Data().(string)
	if !ok {
		name = fmt.Sprintf("Curseforge project %d: %d", projectID, fileID)
	}
	clientOnly, ok := modJson.S("clientOnly").Data().(bool)
	return &CurseForgeModFile{projectID, fileID, name, name, ok && clientOnly}
}

func (f CurseForgeModFile) install(pack *ModPack) error {
	// Check the mod cache to see if we already have the right file ID installed
	lastFileId, lastFilename := pack.modCache.GetLastModFile(f.projectID)
	if lastFileId == f.fileID {
		// Nothing to do; we can skip this installed file
		fmt.Printf("Skipping %s\n", lastFilename)
		return nil
	} else if lastFileId > 0 {
		// A different version of the file is installed; clean it up
		pack.modCache.CleanupModFile(f.projectID)
	}

	// Resolve the project ID into a slug
	slug, err := pack.db.findSlugByProject(f.projectID)
	if err != nil {
		return fmt.Errorf("failed to find slug for project %d: %+v", f.projectID, err)
	}

	// Now, retrieve the JSON descriptor for this file so we can get the CDN url
	descriptorUrl := fmt.Sprintf("https://addons-ecs.forgesvc.net/api/v2/addon/%d/file/%d", f.projectID, f.fileID)
	descriptor, err := getJSONFromURL(descriptorUrl)
	if err != nil {
		return fmt.Errorf("failed to retrieve descriptor for %s: %+v", slug, err)
	}

	// Download the file to the pack mod directory
	finalUrl := descriptor.Path("downloadUrl").Data().(string)

	filename, err := downloadHttpFileToDir(finalUrl, pack.modPath(), true)
	if err != nil {
		return err
	}

	// Download succeeded; register this mod as installed in the cache
	pack.modCache.AddModFile(f.projectID, f.fileID, filename)
	return nil
}

func (f *CurseForgeModFile) update(pack *ModPack) (bool, error) {
	latestFile, err := f.getLatestFile(pack.minecraftVersion())
	if err != nil {
		return false, err
	}

	if latestFile > f.fileID {
		f.fileID = latestFile
		return true, nil
	}

	return false, nil
}

func (f CurseForgeModFile) getName() string {
	return f.name
}

func (f CurseForgeModFile) isClientOnly() bool {
	return f.clientOnly
}

func (f CurseForgeModFile) equalsJson(modJson *gabs.Container) bool {
	projectID, ok := modJson.Path("projectID").Data().(float64)
	return ok && int(projectID) == f.projectID
}

func (f CurseForgeModFile) toJson() map[string]interface{} {
	result := map[string]interface{}{
		"projectID": f.projectID,
		"fileID":    f.fileID,
		"required":  true,
		"desc":      f.name,
	}
	if f.clientOnly {
		result["clientOnly"] = true
	}
	return result
}

func (f CurseForgeModFile) getLatestFile(minecraftVersion string) (int, error) {
	// Pull the project's descriptor, which has a list of the latest files for each version of Minecraft
	projectUrl := fmt.Sprintf("https://addons-ecs.forgesvc.net/api/v2/addon/%d", f.projectID)
	project, err := getJSONFromURL(projectUrl)
	if err != nil {
		return -1, fmt.Errorf("failed to retrieve project for %s: %+v", f.name, err)
	}

	selectedFileType := math.MaxInt8
	selectedFileId := 0

	// Look for the file with the matching version
	files, _ := project.Path("gameVersionLatestFiles").Children()
	for _, file := range files {
		fileType, _ := intValue(file, "fileType") // 1 = release, 2 = beta, 3 = alpha
		fileId, _ := intValue(file, "projectFileId")
		targetVsn := file.Path("gameVersion").Data().(string)

		if targetVsn != minecraftVersion {
			continue
		}

		// Matched on version; prefer releases over beta/alpha
		if fileType < selectedFileType {
			selectedFileType = fileType
			selectedFileId = fileId
		}
	}

	// TODO: Pull file descriptor and check for deps
	return selectedFileId, nil
}
<commit_msg>Return error when no file found for requested minecraft version; thanks TrashboxBobylev<commit_after>// ***************************************************************************
//
//  Copyright 2017 David (Dizzy) Smith, dizzyd@dizzyd.com
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
// ***************************************************************************

package main

import (
	"fmt"
	"math"

	"github.com/Jeffail/gabs"
)

type CurseForgeModFile struct {
	projectID  int
	fileID     int
	desc       string
	name       string
	clientOnly bool
}

func SelectCurseForgeModFile(pack *ModPack, mod string, url string, clientOnly bool) error {
	// Try to find the project ID using the mod name as a slug
	projectID, err := pack.db.findModBySlug(mod)
	if err != nil {
		return fmt.Errorf("unknown mod %s", mod)
	}

	// Look up the slug, name and description
	_, name, desc, err := pack.db.getProjectInfo(projectID)
	if err != nil {
		return fmt.Errorf("no name/description available for %s (%d): %+v", mod, projectID, err)
	}

	// Setup a mod file entry and then pull the latest file info
	modFile := CurseForgeModFile{projectID: projectID, desc: desc, name: name, clientOnly: clientOnly}
	fileId, err := modFile.getLatestFile(pack.minecraftVersion())
	if err != nil {
		return fmt.Errorf("failed to get latest file for %s (%d): %+v", mod, projectID, err)
	}

	// If we found a newer file, update entry and then the pack
	if fileId > modFile.fileID {
		modFile.fileID = fileId
		err = pack.selectMod(&modFile)
		if err != nil {
			return err
		}
	}

	return nil
}

func NewCurseForgeModFile(modJson *gabs.Container) *CurseForgeModFile {
	projectID, _ := intValue(modJson, "projectID")
	fileID, _ := intValue(modJson, "fileID")
	name, ok := modJson.Path("desc").Data().(string)
	if !ok {
		name = fmt.Sprintf("Curseforge project %d: %d", projectID, fileID)
	}
	clientOnly, ok := modJson.S("clientOnly").Data().(bool)
	return &CurseForgeModFile{projectID, fileID, name, name, ok && clientOnly}
}

func (f CurseForgeModFile) install(pack *ModPack) error {
	// Check the mod cache to see if we already have the right file ID installed
	lastFileId, lastFilename := pack.modCache.GetLastModFile(f.projectID)
	if lastFileId == f.fileID {
		// Nothing to do; we can skip this installed file
		fmt.Printf("Skipping %s\n", lastFilename)
		return nil
	} else if lastFileId > 0 {
		// A different version of the file is installed; clean it up
		pack.modCache.CleanupModFile(f.projectID)
	}

	// Resolve the project ID into a slug
	slug, err := pack.db.findSlugByProject(f.projectID)
	if err != nil {
		return fmt.Errorf("failed to find slug for project %d: %+v", f.projectID, err)
	}

	// Now, retrieve the JSON descriptor for this file so we can get the CDN url
	descriptorUrl := fmt.Sprintf("https://addons-ecs.forgesvc.net/api/v2/addon/%d/file/%d", f.projectID, f.fileID)
	descriptor, err := getJSONFromURL(descriptorUrl)
	if err != nil {
		return fmt.Errorf("failed to retrieve descriptor for %s: %+v", slug, err)
	}

	// Download the file to the pack mod directory
	finalUrl := descriptor.Path("downloadUrl").Data().(string)

	filename, err := downloadHttpFileToDir(finalUrl, pack.modPath(), true)
	if err != nil {
		return err
	}

	// Download succeeded; register this mod as installed in the cache
	pack.modCache.AddModFile(f.projectID, f.fileID, filename)
	return nil
}

func (f *CurseForgeModFile) update(pack *ModPack) (bool, error) {
	latestFile, err := f.getLatestFile(pack.minecraftVersion())
	if err != nil {
		return false, err
	}

	if latestFile > f.fileID {
		f.fileID = latestFile
		return true, nil
	}

	return false, nil
}

func (f CurseForgeModFile) getName() string {
	return f.name
}

func (f CurseForgeModFile) isClientOnly() bool {
	return f.clientOnly
}

func (f CurseForgeModFile) equalsJson(modJson *gabs.Container) bool {
	projectID, ok := modJson.Path("projectID").Data().(float64)
	return ok && int(projectID) == f.projectID
}

func (f CurseForgeModFile) toJson() map[string]interface{} {
	result := map[string]interface{}{
		"projectID": f.projectID,
		"fileID":    f.fileID,
		"required":  true,
		"desc":      f.name,
	}
	if f.clientOnly {
		result["clientOnly"] = true
	}
	return result
}

func (f CurseForgeModFile) getLatestFile(minecraftVersion string) (int, error) {
	// Pull the project's descriptor, which has a list of the latest files for each version of Minecraft
	projectUrl := fmt.Sprintf("https://addons-ecs.forgesvc.net/api/v2/addon/%d", f.projectID)
	project, err := getJSONFromURL(projectUrl)
	if err != nil {
		return -1, fmt.Errorf("failed to retrieve project for %s: %+v", f.name, err)
	}

	selectedFileType := math.MaxInt8
	selectedFileId := 0

	// Look for the file with the matching version
	files, _ := project.Path("gameVersionLatestFiles").Children()
	for _, file := range files {
		fileType, _ := intValue(file, "fileType") // 1 = release, 2 = beta, 3 = alpha
		fileId, _ := intValue(file, "projectFileId")
		targetVsn := file.Path("gameVersion").Data().(string)

		if targetVsn != minecraftVersion {
			continue
		}

		// Matched on version; prefer releases over beta/alpha
		if fileType < selectedFileType {
			selectedFileType = fileType
			selectedFileId = fileId
		}
	}

	if selectedFileId == 0 {
		return -1, fmt.Errorf("no version found for Minecraft %s\n", minecraftVersion)
	}

	// TODO: Pull file descriptor and check for deps
	return selectedFileId, nil
}
<|endoftext|>


<commit_before>package connector

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/user"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	"google.golang.org/grpc"

	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dlog"
	"github.com/datawire/telepresence2/pkg/client"
	"github.com/datawire/telepresence2/pkg/client/auth"
	rpc "github.com/datawire/telepresence2/pkg/rpc/connector"
	"github.com/datawire/telepresence2/pkg/rpc/manager"
)

// trafficManager is a handle to access the Traffic Manager in a
// cluster.
type trafficManager struct {
	aiListener      aiListener
	iiListener      iiListener
	conn            *grpc.ClientConn
	grpc            manager.ManagerClient
	startup         chan bool
	apiPort         int32
	sshPort         int32
	userAndHost     string
	installID       string // telepresence's install ID
	sessionID       string // sessionID returned by the traffic-manager
	apiErr          error  // holds the latest traffic-manager API error
	connectCI       bool   // whether --ci was passed to connect
	installer       *installer
	myIntercept     string
	cancelIntercept context.CancelFunc
	// previewHost string // hostname to use for preview URLs, if enabled
}

// newTrafficManager returns a TrafficManager resource for the given
// cluster if it has a Traffic Manager service.
func newTrafficManager(c context.Context, cluster *k8sCluster, installID string, isCI bool) (*trafficManager, error) {
	name, err := user.Current()
	if err != nil {
		return nil, errors.Wrap(err, "user.Current()")
	}
	host, err := os.Hostname()
	if err != nil {
		return nil, errors.Wrap(err, "os.Hostname()")
	}

	// Ensure that we have a traffic-manager to talk to.
	ti, err := newTrafficManagerInstaller(cluster)
	if err != nil {
		return nil, errors.Wrap(err, "new installer")
	}
	localAPIPort, err := getFreePort()
	if err != nil {
		return nil, errors.Wrap(err, "get free port for API")
	}
	localSSHPort, err := getFreePort()
	if err != nil {
		return nil, errors.Wrap(err, "get free port for ssh")
	}
	tm := &trafficManager{
		installer:   ti,
		apiPort:     localAPIPort,
		sshPort:     localSSHPort,
		installID:   installID,
		connectCI:   isCI,
		startup:     make(chan bool),
		userAndHost: fmt.Sprintf("%s@%s", name, host)}

	dgroup.ParentGroup(c).Go("traffic-manager", tm.start)
	return tm, nil
}

func (tm *trafficManager) waitUntilStarted() error {
	<-tm.startup
	return tm.apiErr
}

func (tm *trafficManager) start(c context.Context) error {
	remoteSSHPort, remoteAPIPort, err := tm.installer.ensureManager(c)
	if err != nil {
		tm.apiErr = err
		close(tm.startup)
		return err
	}
	kpfArgs := []string{
		"port-forward",
		"svc/traffic-manager",
		fmt.Sprintf("%d:%d", tm.sshPort, remoteSSHPort),
		fmt.Sprintf("%d:%d", tm.apiPort, remoteAPIPort)}

	err = client.Retry(c, "svc/traffic-manager port-forward", func(c context.Context) error {
		return tm.installer.portForwardAndThen(c, kpfArgs, "init-grpc", tm.initGrpc)
	}, 2*time.Second, 15*time.Second, time.Minute)
	if err != nil && tm.apiErr == nil {
		tm.apiErr = err
		close(tm.startup)
	}
	return err
}

func (tm *trafficManager) bearerToken() string {
	token, err := auth.LoadTokenFromUserCache()
	if err != nil {
		return ""
	}
	return token.AccessToken
}

func (tm *trafficManager) initGrpc(c context.Context) (err error) {
	defer func() {
		tm.apiErr = err
		close(tm.startup)
	}()

	// First check. Establish connection
	tc, cancel := context.WithTimeout(c, connectTimeout)
	defer cancel()

	var conn *grpc.ClientConn
	conn, err = grpc.DialContext(tc, fmt.Sprintf("127.0.0.1:%d", tm.apiPort),
		grpc.WithInsecure(),
		grpc.WithNoProxy(),
		grpc.WithBlock())
	if err != nil {
		if tc.Err() == context.DeadlineExceeded {
			err = errors.New("timeout when connecting to traffic-manager")
		}
		return err
	}

	mClient := manager.NewManagerClient(conn)
	si, err := mClient.ArriveAsClient(c, &manager.ClientInfo{
		Name:        tm.userAndHost,
		InstallId:   tm.installID,
		Product:     "telepresence",
		Version:     client.Version(),
		BearerToken: tm.bearerToken(),
	})

	if err != nil {
		dlog.Errorf(c, "ArriveAsClient: %s", err.Error())
		conn.Close()
		return err
	}
	tm.conn = conn
	tm.grpc = mClient
	tm.sessionID = si.SessionId

	g := dgroup.ParentGroup(c)
	g.Go("remain", tm.remain)
	g.Go("watch-agents", tm.watchAgents)
	g.Go("watch-intercepts", tm.watchIntercepts)
	return nil
}

func (tm *trafficManager) watchAgents(c context.Context) error {
	ac, err := tm.grpc.WatchAgents(c, tm.session())
	if err != nil {
		return err
	}
	return tm.aiListener.start(c, ac)
}

func (tm *trafficManager) watchIntercepts(c context.Context) error {
	ic, err := tm.grpc.WatchIntercepts(c, tm.session())
	if err != nil {
		return err
	}
	return tm.iiListener.start(c, ic)
}

func (tm *trafficManager) session() *manager.SessionInfo {
	return &manager.SessionInfo{SessionId: tm.sessionID}
}

func (tm *trafficManager) agentInfoSnapshot() *manager.AgentInfoSnapshot {
	return tm.aiListener.getData()
}

func (tm *trafficManager) interceptInfoSnapshot() *manager.InterceptInfoSnapshot {
	return tm.iiListener.getData()
}

func (tm *trafficManager) remain(c context.Context) error {
	ticker := time.NewTicker(5 * time.Second)
	for {
		select {
		case <-c.Done():
			return nil
		case <-ticker.C:
			_, err := tm.grpc.Remain(c, &manager.RemainRequest{
				Session:     tm.session(),
				BearerToken: tm.bearerToken(),
			})
			if err != nil {
				return err
			}
		}
	}
}

// Close implements io.Closer
func (tm *trafficManager) Close() error {
	if tm.conn != nil {
		_ = tm.conn.Close()
		tm.conn = nil
		tm.grpc = nil
	}
	return nil
}

func (tm *trafficManager) setStatus(r *rpc.ConnectInfo) {
	if tm.grpc == nil {
		r.Intercepts = &manager.InterceptInfoSnapshot{}
		r.Agents = &manager.AgentInfoSnapshot{}
		if err := tm.apiErr; err != nil {
			r.ErrorText = err.Error()
		}
	} else {
		r.Agents = tm.agentInfoSnapshot()
		r.Intercepts = tm.interceptInfoSnapshot()
	}
}

// A watcher listens on a grpc.ClientStream and notifies listeners when
// something arrives.
type watcher struct {
	entryMaker    func() interface{} // returns an instance of the type produced by the stream
	listeners     []listener
	listenersLock sync.RWMutex
	stream        grpc.ClientStream
}

// watch reads messages from the stream and passes them onto registered listeners. The
// function terminates when the context used when the stream was acquired is cancelled,
// when io.EOF is encountered, or an error occurs during read.
func (r *watcher) watch(c context.Context) error {
	dataChan := make(chan interface{}, 1000)
	defer close(dataChan)

	done := int32(0)
	go func() {
		for {
			select {
			case <-c.Done():
				// ensure no more writes and drain channel to unblock writer
				atomic.StoreInt32(&done, 1)
				for range dataChan {
				}
				return
			case data := <-dataChan:
				if data == nil {
					return
				}

				r.listenersLock.RLock()
				lc := make([]listener, len(r.listeners))
				copy(lc, r.listeners)
				r.listenersLock.RUnlock()

				for _, l := range lc {
					l.onData(data)
				}
			}
		}
	}()

	var err error
	for {
		data := r.entryMaker()
		if err = r.stream.RecvMsg(data); err != nil {
			if err == io.EOF || strings.HasSuffix(err.Error(), " is closing") {
				err = nil
			}
			break
		}
		if atomic.LoadInt32(&done) != 0 {
			break
		}
		dataChan <- data
		if atomic.LoadInt32(&done) != 0 {
			break
		}
	}
	return err
}

func (r *watcher) addListener(l listener) {
	r.listenersLock.Lock()
	r.listeners = append(r.listeners, l)
	r.listenersLock.Unlock()
}

func (r *watcher) removeListener(l listener) {
	r.listenersLock.Lock()
	ls := r.listeners
	for i, x := range ls {
		if l == x {
			last := len(ls) - 1
			ls[i] = ls[last]
			ls[last] = nil
			r.listeners = ls[:last]
			break
		}
	}
	r.listenersLock.Unlock()
}

// A listener gets notified by a watcher when something arrives on the stream
type listener interface {
	onData(data interface{})
}

// An aiListener keeps track of the latest received AgentInfoSnapshot and provides the
// watcher needed to register other listeners.
type aiListener struct {
	watcher
	data atomic.Value
}

func (al *aiListener) getData() *manager.AgentInfoSnapshot {
	v := al.data.Load()
	if v == nil {
		return nil
	}
	return v.(*manager.AgentInfoSnapshot)
}

func (al *aiListener) onData(d interface{}) {
	al.data.Store(d)
}

func (al *aiListener) start(c context.Context, stream grpc.ClientStream) error {
	al.stream = stream
	al.listeners = []listener{al}
	al.entryMaker = func() interface{} { return new(manager.AgentInfoSnapshot) }
	return al.watch(c)
}

func (il *iiListener) onData(d interface{}) {
	il.data.Store(d)
}

func (il *iiListener) start(c context.Context, stream grpc.ClientStream) error {
	il.stream = stream
	il.listeners = []listener{il}
	il.entryMaker = func() interface{} { return new(manager.InterceptInfoSnapshot) }
	return il.watch(c)
}

// iiActive is a listener that waits for an intercept with a given id to become active
type iiActive struct {
	id   string
	done chan *manager.InterceptInfo
}

func (ia *iiActive) onData(d interface{}) {
	if iis, ok := d.(*manager.InterceptInfoSnapshot); ok {
		for _, ii := range iis.Intercepts {
			if ii.Id == ia.id && ii.Disposition != manager.InterceptDispositionType_WAITING {
				done := ia.done
				ia.done = nil
				if done != nil {
					done <- ii
					close(done)
				}
				break
			}
		}
	}
}

// aiPresent is a listener that waits for an agent with a given name to be present
type aiPresent struct {
	name string
	done chan *manager.AgentInfo
}

func (ap *aiPresent) onData(d interface{}) {
	if ais, ok := d.(*manager.AgentInfoSnapshot); ok {
		for _, ai := range ais.Agents {
			if ai.Name == ap.name {
				done := ap.done
				ap.done = nil
				if done != nil {
					done <- ai
					close(done)
				}
				break
			}
		}
	}
}
<commit_msg>connector: Fix getting the local username<commit_after>package connector

import (
	"context"
	"fmt"
	"io"
	"os"
	"os/user"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/pkg/errors"
	"google.golang.org/grpc"

	"github.com/datawire/dlib/dgroup"
	"github.com/datawire/dlib/dlog"
	"github.com/datawire/telepresence2/pkg/client"
	"github.com/datawire/telepresence2/pkg/client/auth"
	rpc "github.com/datawire/telepresence2/pkg/rpc/connector"
	"github.com/datawire/telepresence2/pkg/rpc/manager"
)

// trafficManager is a handle to access the Traffic Manager in a
// cluster.
type trafficManager struct {
	aiListener      aiListener
	iiListener      iiListener
	conn            *grpc.ClientConn
	grpc            manager.ManagerClient
	startup         chan bool
	apiPort         int32
	sshPort         int32
	userAndHost     string
	installID       string // telepresence's install ID
	sessionID       string // sessionID returned by the traffic-manager
	apiErr          error  // holds the latest traffic-manager API error
	connectCI       bool   // whether --ci was passed to connect
	installer       *installer
	myIntercept     string
	cancelIntercept context.CancelFunc
	// previewHost string // hostname to use for preview URLs, if enabled
}

// newTrafficManager returns a TrafficManager resource for the given
// cluster if it has a Traffic Manager service.
func newTrafficManager(c context.Context, cluster *k8sCluster, installID string, isCI bool) (*trafficManager, error) {
	userinfo, err := user.Current()
	if err != nil {
		return nil, errors.Wrap(err, "user.Current()")
	}
	host, err := os.Hostname()
	if err != nil {
		return nil, errors.Wrap(err, "os.Hostname()")
	}

	// Ensure that we have a traffic-manager to talk to.
	ti, err := newTrafficManagerInstaller(cluster)
	if err != nil {
		return nil, errors.Wrap(err, "new installer")
	}
	localAPIPort, err := getFreePort()
	if err != nil {
		return nil, errors.Wrap(err, "get free port for API")
	}
	localSSHPort, err := getFreePort()
	if err != nil {
		return nil, errors.Wrap(err, "get free port for ssh")
	}
	tm := &trafficManager{
		installer:   ti,
		apiPort:     localAPIPort,
		sshPort:     localSSHPort,
		installID:   installID,
		connectCI:   isCI,
		startup:     make(chan bool),
		userAndHost: fmt.Sprintf("%s@%s", userinfo.Username, host)}

	dgroup.ParentGroup(c).Go("traffic-manager", tm.start)
	return tm, nil
}

func (tm *trafficManager) waitUntilStarted() error {
	<-tm.startup
	return tm.apiErr
}

func (tm *trafficManager) start(c context.Context) error {
	remoteSSHPort, remoteAPIPort, err := tm.installer.ensureManager(c)
	if err != nil {
		tm.apiErr = err
		close(tm.startup)
		return err
	}
	kpfArgs := []string{
		"port-forward",
		"svc/traffic-manager",
		fmt.Sprintf("%d:%d", tm.sshPort, remoteSSHPort),
		fmt.Sprintf("%d:%d", tm.apiPort, remoteAPIPort)}

	err = client.Retry(c, "svc/traffic-manager port-forward", func(c context.Context) error {
		return tm.installer.portForwardAndThen(c, kpfArgs, "init-grpc", tm.initGrpc)
	}, 2*time.Second, 15*time.Second, time.Minute)
	if err != nil && tm.apiErr == nil {
		tm.apiErr = err
		close(tm.startup)
	}
	return err
}

func (tm *trafficManager) bearerToken() string {
	token, err := auth.LoadTokenFromUserCache()
	if err != nil {
		return ""
	}
	return token.AccessToken
}

func (tm *trafficManager) initGrpc(c context.Context) (err error) {
	defer func() {
		tm.apiErr = err
		close(tm.startup)
	}()

	// First check. Establish connection
	tc, cancel := context.WithTimeout(c, connectTimeout)
	defer cancel()

	var conn *grpc.ClientConn
	conn, err = grpc.DialContext(tc, fmt.Sprintf("127.0.0.1:%d", tm.apiPort),
		grpc.WithInsecure(),
		grpc.WithNoProxy(),
		grpc.WithBlock())
	if err != nil {
		if tc.Err() == context.DeadlineExceeded {
			err = errors.New("timeout when connecting to traffic-manager")
		}
		return err
	}

	mClient := manager.NewManagerClient(conn)
	si, err := mClient.ArriveAsClient(c, &manager.ClientInfo{
		Name:        tm.userAndHost,
		InstallId:   tm.installID,
		Product:     "telepresence",
		Version:     client.Version(),
		BearerToken: tm.bearerToken(),
	})

	if err != nil {
		dlog.Errorf(c, "ArriveAsClient: %s", err.Error())
		conn.Close()
		return err
	}
	tm.conn = conn
	tm.grpc = mClient
	tm.sessionID = si.SessionId

	g := dgroup.ParentGroup(c)
	g.Go("remain", tm.remain)
	g.Go("watch-agents", tm.watchAgents)
	g.Go("watch-intercepts", tm.watchIntercepts)
	return nil
}

func (tm *trafficManager) watchAgents(c context.Context) error {
	ac, err := tm.grpc.WatchAgents(c, tm.session())
	if err != nil {
		return err
	}
	return tm.aiListener.start(c, ac)
}

func (tm *trafficManager) watchIntercepts(c context.Context) error {
	ic, err := tm.grpc.WatchIntercepts(c, tm.session())
	if err != nil {
		return err
	}
	return tm.iiListener.start(c, ic)
}

func (tm *trafficManager) session() *manager.SessionInfo {
	return &manager.SessionInfo{SessionId: tm.sessionID}
}

func (tm *trafficManager) agentInfoSnapshot() *manager.AgentInfoSnapshot {
	return tm.aiListener.getData()
}

func (tm *trafficManager) interceptInfoSnapshot() *manager.InterceptInfoSnapshot {
	return tm.iiListener.getData()
}

func (tm *trafficManager) remain(c context.Context) error {
	ticker := time.NewTicker(5 * time.Second)
	for {
		select {
		case <-c.Done():
			return nil
		case <-ticker.C:
			_, err := tm.grpc.Remain(c, &manager.RemainRequest{
				Session:     tm.session(),
				BearerToken: tm.bearerToken(),
			})
			if err != nil {
				return err
			}
		}
	}
}

// Close implements io.Closer
func (tm *trafficManager) Close() error {
	if tm.conn != nil {
		_ = tm.conn.Close()
		tm.conn = nil
		tm.grpc = nil
	}
	return nil
}

func (tm *trafficManager) setStatus(r *rpc.ConnectInfo) {
	if tm.grpc == nil {
		r.Intercepts = &manager.InterceptInfoSnapshot{}
		r.Agents = &manager.AgentInfoSnapshot{}
		if err := tm.apiErr; err != nil {
			r.ErrorText = err.Error()
		}
	} else {
		r.Agents = tm.agentInfoSnapshot()
		r.Intercepts = tm.interceptInfoSnapshot()
	}
}

// A watcher listens on a grpc.ClientStream and notifies listeners when
// something arrives.
type watcher struct {
	entryMaker    func() interface{} // returns an instance of the type produced by the stream
	listeners     []listener
	listenersLock sync.RWMutex
	stream        grpc.ClientStream
}

// watch reads messages from the stream and passes them onto registered listeners. The
// function terminates when the context used when the stream was acquired is cancelled,
// when io.EOF is encountered, or an error occurs during read.
func (r *watcher) watch(c context.Context) error {
	dataChan := make(chan interface{}, 1000)
	defer close(dataChan)

	done := int32(0)
	go func() {
		for {
			select {
			case <-c.Done():
				// ensure no more writes and drain channel to unblock writer
				atomic.StoreInt32(&done, 1)
				for range dataChan {
				}
				return
			case data := <-dataChan:
				if data == nil {
					return
				}

				r.listenersLock.RLock()
				lc := make([]listener, len(r.listeners))
				copy(lc, r.listeners)
				r.listenersLock.RUnlock()

				for _, l := range lc {
					l.onData(data)
				}
			}
		}
	}()

	var err error
	for {
		data := r.entryMaker()
		if err = r.stream.RecvMsg(data); err != nil {
			if err == io.EOF || strings.HasSuffix(err.Error(), " is closing") {
				err = nil
			}
			break
		}
		if atomic.LoadInt32(&done) != 0 {
			break
		}
		dataChan <- data
		if atomic.LoadInt32(&done) != 0 {
			break
		}
	}
	return err
}

func (r *watcher) addListener(l listener) {
	r.listenersLock.Lock()
	r.listeners = append(r.listeners, l)
	r.listenersLock.Unlock()
}

func (r *watcher) removeListener(l listener) {
	r.listenersLock.Lock()
	ls := r.listeners
	for i, x := range ls {
		if l == x {
			last := len(ls) - 1
			ls[i] = ls[last]
			ls[last] = nil
			r.listeners = ls[:last]
			break
		}
	}
	r.listenersLock.Unlock()
}

// A listener gets notified by a watcher when something arrives on the stream
type listener interface {
	onData(data interface{})
}

// An aiListener keeps track of the latest received AgentInfoSnapshot and provides the
// watcher needed to register other listeners.
type aiListener struct {
	watcher
	data atomic.Value
}

func (al *aiListener) getData() *manager.AgentInfoSnapshot {
	v := al.data.Load()
	if v == nil {
		return nil
	}
	return v.(*manager.AgentInfoSnapshot)
}

func (al *aiListener) onData(d interface{}) {
	al.data.Store(d)
}

func (al *aiListener) start(c context.Context, stream grpc.ClientStream) error {
	al.stream = stream
	al.listeners = []listener{al}
	al.entryMaker = func() interface{} { return new(manager.AgentInfoSnapshot) }
	return al.watch(c)
}

func (il *iiListener) onData(d interface{}) {
	il.data.Store(d)
}

func (il *iiListener) start(c context.Context, stream grpc.ClientStream) error {
	il.stream = stream
	il.listeners = []listener{il}
	il.entryMaker = func() interface{} { return new(manager.InterceptInfoSnapshot) }
	return il.watch(c)
}

// iiActive is a listener that waits for an intercept with a given id to become active
type iiActive struct {
	id   string
	done chan *manager.InterceptInfo
}

func (ia *iiActive) onData(d interface{}) {
	if iis, ok := d.(*manager.InterceptInfoSnapshot); ok {
		for _, ii := range iis.Intercepts {
			if ii.Id == ia.id && ii.Disposition != manager.InterceptDispositionType_WAITING {
				done := ia.done
				ia.done = nil
				if done != nil {
					done <- ii
					close(done)
				}
				break
			}
		}
	}
}

// aiPresent is a listener that waits for an agent with a given name to be present
type aiPresent struct {
	name string
	done chan *manager.AgentInfo
}

func (ap *aiPresent) onData(d interface{}) {
	if ais, ok := d.(*manager.AgentInfoSnapshot); ok {
		for _, ai := range ais.Agents {
			if ai.Name == ap.name {
				done := ap.done
				ap.done = nil
				if done != nil {
					done <- ai
					close(done)
				}
				break
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The etcd-operator 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 controller

import (
	api "github.com/coreos/etcd-operator/pkg/apis/etcd/v1beta2"
	"github.com/coreos/etcd-operator/pkg/backup/backupapi"
	"github.com/coreos/etcd-operator/pkg/util/etcdutil"
	"github.com/coreos/etcd-operator/pkg/util/k8sutil"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

const (
	// Copy from deployment_controller.go:
	// maxRetries is the number of times a restore request will be retried before it is dropped out of the queue.
	// With the current rate-limiter in use (5ms*2^(maxRetries-1)) the following numbers represent the times
	// an restore request is going to be requeued:
	//
	// 5ms, 10ms, 20ms, 40ms, 80ms, 160ms, 320ms, 640ms, 1.3s, 2.6s, 5.1s, 10.2s, 20.4s, 41s, 82s
	maxRetries = 15
)

func (r *Restore) runWorker() {
	for r.processNextItem() {
	}
}

func (r *Restore) processNextItem() bool {
	// Wait until there is a new item in the working queue
	key, quit := r.queue.Get()
	if quit {
		return false
	}
	// Tell the queue that we are done with processing this key. This unblocks the key for other workers
	// This allows safe parallel processing because two pods with the same key are never processed in
	// parallel.
	defer r.queue.Done(key)
	err := r.processItem(key.(string))
	// Handle the error if something went wrong during the execution of the business logic
	r.handleErr(err, key)
	return true
}

func (r *Restore) processItem(key string) error {
	obj, exists, err := r.indexer.GetByKey(key)
	if err != nil {
		return err
	}
	if !exists {
		cn, ok := r.clusterNames.Load(key)
		if ok {
			r.restoreCRs.Delete(cn)
			r.clusterNames.Delete(key)
		}
		return nil
	}

	er := obj.(*api.EtcdRestore)
	clusterName := er.Spec.BackupSpec.ClusterName
	r.clusterNames.Store(key, clusterName)
	r.restoreCRs.Store(clusterName, er)
	return createSeedPod(r.kubecli, er.Spec.ClusterSpec, er.AsOwner(), r.namespace, er.Spec.ClusterSpec.Version, r.mySvcAddr, clusterName)
}

func createSeedPod(kubecli kubernetes.Interface, cs api.ClusterSpec, owner metav1.OwnerReference, namespace, etcdVersion, svcAddr, clusterName string) error {
	err := k8sutil.CreateClientService(kubecli, clusterName, namespace, owner)
	if err != nil {
		return err
	}
	err = k8sutil.CreatePeerService(kubecli, clusterName, namespace, owner)
	if err != nil {
		return err
	}

	m := &etcdutil.Member{
		Name:      etcdutil.CreateMemberName(clusterName, 0),
		Namespace: namespace,
		// TODO: support TLS
		SecurePeer:   false,
		SecureClient: false,
	}
	ms := etcdutil.NewMemberSet(m)
	backupURL := backupapi.BackupURLForCluster("http", svcAddr, clusterName, etcdVersion, -1)
	cs.Cleanup()
	pod := k8sutil.NewSeedMemberPod(clusterName, ms, m, cs, owner, backupURL)
	_, err = kubecli.Core().Pods(namespace).Create(pod)
	return err
}

func (r *Restore) handleErr(err error, key interface{}) {
	if err == nil {
		// Forget about the #AddRateLimited history of the key on every successful synchronization.
		// This ensures that future processing of updates for this key is not delayed because of
		// an outdated error history.
		r.queue.Forget(key)
		return
	}

	// This controller retries maxRetries times if something goes wrong. After that, it stops trying.
	if r.queue.NumRequeues(key) < maxRetries {
		r.logger.Errorf("error syncing restore request (%v): %v", key, err)

		// Re-enqueue the key rate limited. Based on the rate limiter on the
		// queue and the re-enqueue history, the key will be processed later again.
		r.queue.AddRateLimited(key)
		return
	}

	r.queue.Forget(key)
	// Report that, even after several retries, we could not successfully process this key
	r.logger.Infof("dropping restore request (%v) out of the queue: %v", key, err)
}
<commit_msg>contoller/restore-operator: report restore status<commit_after>// Copyright 2017 The etcd-operator 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 controller

import (
	api "github.com/coreos/etcd-operator/pkg/apis/etcd/v1beta2"
	"github.com/coreos/etcd-operator/pkg/backup/backupapi"
	"github.com/coreos/etcd-operator/pkg/util/etcdutil"
	"github.com/coreos/etcd-operator/pkg/util/k8sutil"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

const (
	// Copy from deployment_controller.go:
	// maxRetries is the number of times a restore request will be retried before it is dropped out of the queue.
	// With the current rate-limiter in use (5ms*2^(maxRetries-1)) the following numbers represent the times
	// an restore request is going to be requeued:
	//
	// 5ms, 10ms, 20ms, 40ms, 80ms, 160ms, 320ms, 640ms, 1.3s, 2.6s, 5.1s, 10.2s, 20.4s, 41s, 82s
	maxRetries = 15
)

func (r *Restore) runWorker() {
	for r.processNextItem() {
	}
}

func (r *Restore) processNextItem() bool {
	// Wait until there is a new item in the working queue
	key, quit := r.queue.Get()
	if quit {
		return false
	}
	// Tell the queue that we are done with processing this key. This unblocks the key for other workers
	// This allows safe parallel processing because two pods with the same key are never processed in
	// parallel.
	defer r.queue.Done(key)
	err := r.processItem(key.(string))
	// Handle the error if something went wrong during the execution of the business logic
	r.handleErr(err, key)
	return true
}

func (r *Restore) processItem(key string) error {
	obj, exists, err := r.indexer.GetByKey(key)
	if err != nil {
		return err
	}
	if !exists {
		cn, ok := r.clusterNames.Load(key)
		if ok {
			r.restoreCRs.Delete(cn)
			r.clusterNames.Delete(key)
		}
		return nil
	}

	er := obj.(*api.EtcdRestore)
	clusterName := er.Spec.BackupSpec.ClusterName
	r.clusterNames.Store(key, clusterName)
	r.restoreCRs.Store(clusterName, er)
	err = createSeedPod(r.kubecli, er.Spec.ClusterSpec, er.AsOwner(), r.namespace, er.Spec.ClusterSpec.Version, r.mySvcAddr, clusterName)
	r.reportStatus(err, er)
	return err
}

func (r *Restore) reportStatus(err error, er *api.EtcdRestore) {
	if err != nil {
		er.Status.Succeeded = false
		er.Status.Reason = err.Error()
	} else {
		er.Status.Succeeded = true
	}
	_, err = r.restoreCRCli.EtcdV1beta2().EtcdRestores(r.namespace).Update(er)
	if err != nil {
		r.logger.Warningf("failed to update status of restore CR %v : (%v)", er.Name, err)
	}
}

func createSeedPod(kubecli kubernetes.Interface, cs api.ClusterSpec, owner metav1.OwnerReference, namespace, etcdVersion, svcAddr, clusterName string) error {
	err := k8sutil.CreateClientService(kubecli, clusterName, namespace, owner)
	if err != nil {
		return err
	}
	err = k8sutil.CreatePeerService(kubecli, clusterName, namespace, owner)
	if err != nil {
		return err
	}

	m := &etcdutil.Member{
		Name:      etcdutil.CreateMemberName(clusterName, 0),
		Namespace: namespace,
		// TODO: support TLS
		SecurePeer:   false,
		SecureClient: false,
	}
	ms := etcdutil.NewMemberSet(m)
	backupURL := backupapi.BackupURLForCluster("http", svcAddr, clusterName, etcdVersion, -1)
	cs.Cleanup()
	pod := k8sutil.NewSeedMemberPod(clusterName, ms, m, cs, owner, backupURL)
	_, err = kubecli.Core().Pods(namespace).Create(pod)
	return err
}

func (r *Restore) handleErr(err error, key interface{}) {
	if err == nil {
		// Forget about the #AddRateLimited history of the key on every successful synchronization.
		// This ensures that future processing of updates for this key is not delayed because of
		// an outdated error history.
		r.queue.Forget(key)
		return
	}

	// This controller retries maxRetries times if something goes wrong. After that, it stops trying.
	if r.queue.NumRequeues(key) < maxRetries {
		r.logger.Errorf("error syncing restore request (%v): %v", key, err)

		// Re-enqueue the key rate limited. Based on the rate limiter on the
		// queue and the re-enqueue history, the key will be processed later again.
		r.queue.AddRateLimited(key)
		return
	}

	r.queue.Forget(key)
	// Report that, even after several retries, we could not successfully process this key
	r.logger.Infof("dropping restore request (%v) out of the queue: %v", key, err)
}
<|endoftext|>


<commit_before>package payment_method

import (
	"database/sql"
	"testing"

	"github.com/fritzpay/paymentd/pkg/paymentd/principal"
	"github.com/fritzpay/paymentd/pkg/paymentd/project"
	"github.com/fritzpay/paymentd/pkg/paymentd/provider"
	"github.com/fritzpay/paymentd/pkg/testutil"
	. "github.com/smartystreets/goconvey/convey"
)

func TestPaymentMethodSQL(t *testing.T) {
	Convey("Given a payment DB connection", t, testutil.WithPaymentDB(t, func(db *sql.DB) {
		Reset(func() {
			db.Close()
		})
		Convey("Given a principal DB connection", testutil.WithPrincipalDB(t, func(prDB *sql.DB) {
			Reset(func() {
				prDB.Close()
			})
			Convey("Given a test principal", func() {
				princ, err := principal.PrincipalByNameDB(prDB, "testprincipal")
				So(err, ShouldBeNil)
				So(princ.ID, ShouldNotEqual, 0)
				So(princ.Empty(), ShouldBeFalse)

				Convey("Given a test project", func() {
					proj, err := project.ProjectByPrincipalIDNameDB(prDB, princ.ID, "testproject")
					So(err, ShouldBeNil)

					Convey("Given a transaction", func() {
						tx, err := db.Begin()
						So(err, ShouldBeNil)

						Reset(func() {
							err = tx.Rollback()
							So(err, ShouldBeNil)
						})

						Convey("Given a test provider exists", func() {
							pr, err := provider.ProviderByIDTx(tx, 1)
							So(err, ShouldBeNil)
							So(pr.ID, ShouldEqual, 1)

							Convey("When retrieving a nonexistent payment method", func() {
								_, err = PaymentMethodByProjectIDProviderIDMethodKeyDB(db, proj.ID, pr.ID, "test")
								Convey("It should return a not found error", func() {
									So(err, ShouldEqual, ErrPaymentMethodNotFound)
								})
							})

							Convey("When inserting a new payment method", func() {
								pm := &Method{}
								pm.ProjectID = proj.ID
								pm.Provider.ID = pr.ID
								pm.MethodKey = "testInsert"
								pm.CreatedBy = "test"

								err = InsertPaymentMethodTx(tx, pm)
								So(err, ShouldBeNil)
								So(pm.ID, ShouldNotEqual, 0)

								Convey("When setting the status to active", func() {
									pm.Status = PaymentMethodStatusActive
									pm.CreatedBy = "test"

									err = InsertPaymentMethodStatusTx(tx, pm)
									So(err, ShouldBeNil)

									Convey("When retrieving the payment method", func() {
										newPm, err := PaymentMethodByIDTx(tx, pm.ID)
										So(err, ShouldBeNil)

										Convey("The retrieved payment method should match", func() {
											So(newPm.Status, ShouldEqual, PaymentMethodStatusActive)
										})
									})
								})

								Convey("When setting metadata", func() {
									pm.Metadata = map[string]string{
										"name": "value",
										"test": "check",
									}
									err = InsertPaymentMethodMetadataTx(tx, pm, "metatest")
									So(err, ShouldBeNil)

									Convey("When selecting metadata", func() {
										metadata, err := PaymentMethodMetadataTx(tx, pm)
										So(err, ShouldBeNil)

										Convey("It should match", func() {
											So(metadata, ShouldNotBeNil)
											So(metadata["name"], ShouldEqual, "value")
											So(metadata["test"], ShouldEqual, "check")
										})
									})
								})
							})
						})
					})
				})
			})
		}))
	}))
}
<commit_msg>fixed test<commit_after>package payment_method

import (
	"database/sql"
	"testing"

	"github.com/fritzpay/paymentd/pkg/paymentd/principal"
	"github.com/fritzpay/paymentd/pkg/paymentd/project"
	"github.com/fritzpay/paymentd/pkg/paymentd/provider"
	"github.com/fritzpay/paymentd/pkg/testutil"
	. "github.com/smartystreets/goconvey/convey"
)

func TestPaymentMethodSQL(t *testing.T) {
	Convey("Given a payment DB connection", t, testutil.WithPaymentDB(t, func(db *sql.DB) {
		Reset(func() {
			db.Close()
		})
		Convey("Given a principal DB connection", testutil.WithPrincipalDB(t, func(prDB *sql.DB) {
			Reset(func() {
				prDB.Close()
			})
			Convey("Given a test principal", func() {
				princ, err := principal.PrincipalByNameDB(prDB, "testprincipal")
				So(err, ShouldBeNil)
				So(princ.ID, ShouldNotEqual, 0)
				So(princ.Empty(), ShouldBeFalse)

				Convey("Given a test project", func() {
					proj, err := project.ProjectByPrincipalIDNameDB(prDB, princ.ID, "testproject")
					So(err, ShouldBeNil)

					Convey("Given a transaction", func() {
						tx, err := db.Begin()
						So(err, ShouldBeNil)

						Reset(func() {
							err = tx.Rollback()
							So(err, ShouldBeNil)
						})

						Convey("Given a test provider exists", func() {
							pr, err := provider.ProviderByIDTx(tx, 1)
							So(err, ShouldBeNil)
							So(pr.ID, ShouldEqual, 1)

							Convey("When retrieving a nonexistent payment method", func() {
								_, err = PaymentMethodByProjectIDProviderIDMethodKeyTx(tx, proj.ID, pr.ID, "nonexistent")
								Convey("It should return a not found error", func() {
									So(err, ShouldEqual, ErrPaymentMethodNotFound)
								})
							})

							Convey("When retrieving an existent payment method", func() {
								pm, err := PaymentMethodByProjectIDProviderIDMethodKeyTx(tx, proj.ID, pr.ID, "test")
								Convey("It should return a payment method", func() {
									So(err, ShouldBeNil)
									So(pm.MethodKey, ShouldEqual, "test")
								})
							})

							Convey("When inserting a new payment method", func() {
								pm := &Method{}
								pm.ProjectID = proj.ID
								pm.Provider.ID = pr.ID
								pm.MethodKey = "testInsert"
								pm.CreatedBy = "test"

								err = InsertPaymentMethodTx(tx, pm)
								So(err, ShouldBeNil)
								So(pm.ID, ShouldNotEqual, 0)

								Convey("When setting the status to active", func() {
									pm.Status = PaymentMethodStatusActive
									pm.CreatedBy = "test"

									err = InsertPaymentMethodStatusTx(tx, pm)
									So(err, ShouldBeNil)

									Convey("When retrieving the payment method", func() {
										newPm, err := PaymentMethodByIDTx(tx, pm.ID)
										So(err, ShouldBeNil)

										Convey("The retrieved payment method should match", func() {
											So(newPm.Status, ShouldEqual, PaymentMethodStatusActive)
										})
									})
								})

								Convey("When setting metadata", func() {
									pm.Metadata = map[string]string{
										"name": "value",
										"test": "check",
									}
									err = InsertPaymentMethodMetadataTx(tx, pm, "metatest")
									So(err, ShouldBeNil)

									Convey("When selecting metadata", func() {
										metadata, err := PaymentMethodMetadataTx(tx, pm)
										So(err, ShouldBeNil)

										Convey("It should match", func() {
											So(metadata, ShouldNotBeNil)
											So(metadata["name"], ShouldEqual, "value")
											So(metadata["test"], ShouldEqual, "check")
										})
									})
								})
							})
						})
					})
				})
			})
		}))
	}))
}
<|endoftext|>


<commit_before>/*
Copyright 2014 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 registrytest

import (
	"sync"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/watch"
)

// TODO: Why do we have this AND MemoryRegistry?
type ControllerRegistry struct {
	Err         error
	Controllers *api.ReplicationControllerList
	sync.Mutex
}

func (r *ControllerRegistry) SetError(err error) {
	r.Lock()
	defer r.Unlock()
	r.Err = err
}

func (r *ControllerRegistry) ListControllers(ctx api.Context) (*api.ReplicationControllerList, error) {
	return r.Controllers, r.Err
}

func (r *ControllerRegistry) GetController(ctx api.Context, ID string) (*api.ReplicationController, error) {
	return &api.ReplicationController{}, r.Err
}

func (r *ControllerRegistry) CreateController(ctx api.Context, controller *api.ReplicationController) error {
	r.Lock()
	defer r.Unlock()
	return r.Err
}

func (r *ControllerRegistry) UpdateController(ctx api.Context, controller *api.ReplicationController) error {
	return r.Err
}

func (r *ControllerRegistry) DeleteController(ctx api.Context, ID string) error {
	return r.Err
}

func (r *ControllerRegistry) WatchControllers(ctx api.Context, label, field labels.Selector, resourceVersion string) (watch.Interface, error) {
	return nil, r.Err
}
<commit_msg>Add some locking to prevent go data race logic from tripping.<commit_after>/*
Copyright 2014 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 registrytest

import (
	"sync"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/watch"
)

// TODO: Why do we have this AND MemoryRegistry?
type ControllerRegistry struct {
	Err         error
	Controllers *api.ReplicationControllerList
	sync.Mutex
}

func (r *ControllerRegistry) SetError(err error) {
	r.Lock()
	defer r.Unlock()
	r.Err = err
}

func (r *ControllerRegistry) ListControllers(ctx api.Context) (*api.ReplicationControllerList, error) {
	r.Lock()
	defer r.Unlock()
	return r.Controllers, r.Err
}

func (r *ControllerRegistry) GetController(ctx api.Context, ID string) (*api.ReplicationController, error) {
	r.Lock()
	defer r.Unlock()
	return &api.ReplicationController{}, r.Err
}

func (r *ControllerRegistry) CreateController(ctx api.Context, controller *api.ReplicationController) error {
	r.Lock()
	defer r.Unlock()
	return r.Err
}

func (r *ControllerRegistry) UpdateController(ctx api.Context, controller *api.ReplicationController) error {
	r.Lock()
	defer r.Unlock()
	return r.Err
}

func (r *ControllerRegistry) DeleteController(ctx api.Context, ID string) error {
	r.Lock()
	defer r.Unlock()
	return r.Err
}

func (r *ControllerRegistry) WatchControllers(ctx api.Context, label, field labels.Selector, resourceVersion string) (watch.Interface, error) {
	r.Lock()
	defer r.Unlock()
	return nil, r.Err
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google LLC

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 build

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"path/filepath"
	"time"

	cstorage "cloud.google.com/go/storage"
	"golang.org/x/oauth2/google"
	cloudbuild "google.golang.org/api/cloudbuild/v1"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/iterator"

	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/build/tag"
	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/config"
	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/constants"
	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/docker"
	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

const (
	// StatusUnknown "STATUS_UNKNOWN" - Status of the build is unknown.
	StatusUnknown = "STATUS_UNKNOWN"

	// StatusQueued "QUEUED" - Build is queued; work has not yet begun.
	StatusQueued = "QUEUED"

	// StatusWorking "WORKING" - Build is being executed.
	StatusWorking = "WORKING"

	// StatusSuccess  "SUCCESS" - Build finished successfully.
	StatusSuccess = "SUCCESS"

	// StatusFailure  "FAILURE" - Build failed to complete successfully.
	StatusFailure = "FAILURE"

	// StatusInternalError  "INTERNAL_ERROR" - Build failed due to an internal cause.
	StatusInternalError = "INTERNAL_ERROR"

	// StatusTimeout  "TIMEOUT" - Build took longer than was allowed.
	StatusTimeout = "TIMEOUT"

	// StatusCancelled  "CANCELLED" - Build was canceled by a user.
	StatusCancelled = "CANCELLED"

	// RetryDelay is the time to wait in between polling the status of the cloud build
	RetryDelay = 1 * time.Second
)

type GoogleCloudBuilder struct {
	*config.BuildConfig
}

func NewGoogleCloudBuilder(cfg *config.BuildConfig) (*GoogleCloudBuilder, error) {
	return &GoogleCloudBuilder{cfg}, nil
}

func (cb *GoogleCloudBuilder) Build(ctx context.Context, out io.Writer, tagger tag.Tagger, artifacts []*config.Artifact) (*BuildResult, error) {
	client, err := google.DefaultClient(ctx, cloudbuild.CloudPlatformScope)
	if err != nil {
		return nil, errors.Wrap(err, "getting google client")
	}
	cbclient, err := cloudbuild.New(client)
	if err != nil {
		return nil, errors.Wrap(err, "getting builder")
	}
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting cloud storage client")
	}
	defer c.Close()
	builds := []Build{}
	for _, artifact := range artifacts {
		build, err := cb.buildArtifact(ctx, out, cbclient, c, artifact)
		if err != nil {
			return nil, errors.Wrapf(err, "building artifact %s", artifact.ImageName)
		}
		builds = append(builds, *build)
	}

	return &BuildResult{
		Builds: builds,
	}, nil
}

func (cb *GoogleCloudBuilder) buildArtifact(ctx context.Context, out io.Writer, cbclient *cloudbuild.Service, c *cstorage.Client, artifact *config.Artifact) (*Build, error) {
	if artifact.DockerArtifact == nil {
		artifact.DockerArtifact = config.DefaultDockerArtifact
	}
	logrus.Infof("Building artifact: %+v", artifact)
	cbBucket := fmt.Sprintf("%s%s", cb.GoogleCloudBuild.ProjectID, constants.GCSBucketSuffix)
	buildObject := fmt.Sprintf("source/%s-%s.tar.gz", cb.GoogleCloudBuild.ProjectID, util.RandomID())

	if err := cb.createBucketIfNotExists(ctx, cbBucket); err != nil {
		return nil, errors.Wrap(err, "creating bucket if not exists")
	}
	if err := cb.checkBucketProjectCorrect(ctx, cbBucket); err != nil {
		return nil, errors.Wrap(err, "checking bucket is in correct project")
	}

	fmt.Fprintf(out, "Pushing code to gs://%s/%s\n", cbBucket, buildObject)
	if err := cb.uploadTarToGCS(ctx, artifact.DockerArtifact.DockerfilePath, artifact.Workspace, cbBucket, buildObject); err != nil {
		return nil, errors.Wrap(err, "uploading source tarball")
	}
	call := cbclient.Projects.Builds.Create(cb.GoogleCloudBuild.ProjectID, &cloudbuild.Build{
		LogsBucket: cbBucket,
		Source: &cloudbuild.Source{
			StorageSource: &cloudbuild.StorageSource{
				Bucket: cbBucket,
				Object: buildObject,
			},
		},
		Steps: []*cloudbuild.BuildStep{
			{
				Name: "gcr.io/cloud-builders/docker",
				Args: []string{"build", "--tag", artifact.ImageName, "-f", artifact.DockerArtifact.DockerfilePath, "."},
			},
		},
		Images: []string{artifact.ImageName},
	})
	op, err := call.Context(ctx).Do()
	if err != nil {
		return nil, errors.Wrap(err, "could not create build")
	}

	remoteID, err := getBuildID(op)
	if err != nil {
		return nil, errors.Wrapf(err, "getting build ID from op")
	}
	logsObject := fmt.Sprintf("log-%s.txt", remoteID)
	fmt.Fprintf(out, "Logs at available at \nhttps://console.cloud.google.com/m/cloudstorage/b/%s/o/%s\n", cbBucket, logsObject)
	var imageID string
	offset := int64(0)
watch:
	for {
		logrus.Debugf("current offset %d", offset)
		b, err := cbclient.Projects.Builds.Get(cb.GoogleCloudBuild.ProjectID, remoteID).Do()
		if err != nil {
			return nil, errors.Wrap(err, "getting build status")
		}

		r, err := cb.getLogs(ctx, offset, cbBucket, logsObject)
		if err != nil {
			return nil, errors.Wrap(err, "getting logs")
		}
		if r != nil {
			written, err := io.Copy(out, r)
			if err != nil {
				return nil, errors.Wrap(err, "copying logs to stdout")
			}
			offset += written
			r.Close()
		}
		switch b.Status {
		case StatusQueued, StatusWorking, StatusUnknown:
		case StatusSuccess:
			imageID, err = getImageID(b)
			if err != nil {
				return nil, errors.Wrap(err, "getting image id from finished build")
			}
			break watch
		case StatusFailure, StatusInternalError, StatusTimeout, StatusCancelled:
			return nil, fmt.Errorf("cloud build failed: %s", b.Status)
		default:
			return nil, fmt.Errorf("unknown status: %s", b.Status)
		}

		time.Sleep(RetryDelay)
	}

	if err := c.Bucket(cbBucket).Object(buildObject).Delete(ctx); err != nil {
		return nil, errors.Wrap(err, "cleaning up source tar after build")
	}
	logrus.Infof("Deleted object %s", buildObject)
	tag := fmt.Sprintf("%s@%s", artifact.ImageName, imageID)
	logrus.Infof("Image built at %s", tag)
	return &Build{
		ImageName: artifact.ImageName,
		Tag:       tag,
		Artifact:  artifact,
	}, nil
}

func getBuildID(op *cloudbuild.Operation) (string, error) {
	if op.Metadata == nil {
		return "", errors.New("missing Metadata in operation")
	}
	var buildMeta cloudbuild.BuildOperationMetadata
	if err := json.Unmarshal([]byte(op.Metadata), &buildMeta); err != nil {
		return "", err
	}
	if buildMeta.Build == nil {
		return "", errors.New("missing Build in operation metadata")
	}
	return buildMeta.Build.Id, nil
}

func getImageID(b *cloudbuild.Build) (string, error) {
	if b.Results == nil || len(b.Results.Images) == 0 {
		return "", errors.New("build failed")
	}
	return b.Results.Images[0].Digest, nil
}

func (cb *GoogleCloudBuilder) uploadTarToGCS(ctx context.Context, dockerfilePath, dockerCtx, bucket, objectName string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return err
	}
	defer c.Close()

	relDockerfilePath := filepath.Join(dockerCtx, dockerfilePath)
	w := c.Bucket(bucket).Object(objectName).NewWriter(ctx)
	if err := docker.CreateDockerTarGzContext(w, relDockerfilePath, dockerCtx); err != nil {
		return errors.Wrap(err, "uploading targz to google storage")
	}
	return w.Close()
}

func (cb *GoogleCloudBuilder) getLogs(ctx context.Context, offset int64, bucket, objectName string) (io.ReadCloser, error) {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting storage client")
	}
	defer c.Close()

	r, err := c.Bucket(bucket).Object(objectName).NewRangeReader(ctx, offset, -1)
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			switch gerr.Code {
			case 404, 416, 429, 503:
				logrus.Debugf("Status Code: %d, %s", gerr.Code, gerr.Body)
				return nil, nil
			}
		}
		if err == cstorage.ErrObjectNotExist {
			logrus.Debugf("Logs for %s %s not uploaded yet...", bucket, objectName)
			return nil, nil
		}
		return nil, errors.Wrap(err, "unknown error")
	}
	return r, nil
}

func (cb *GoogleCloudBuilder) checkBucketProjectCorrect(ctx context.Context, bucket string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return errors.Wrap(err, "getting storage client")
	}
	it := c.Buckets(ctx, cb.GoogleCloudBuild.ProjectID)
	// Set the prefix to the bucket we're looking for to only return that bucket and buckets with that prefix
	// that we'll filter further later on
	it.Prefix = bucket
	for {
		attrs, err := it.Next()
		if err == iterator.Done {
			return errors.Wrap(err, "bucket not found")
		}
		if err != nil {
			return errors.Wrap(err, "iterating over buckets")
		}
		// Since we can't filter on bucket name specifically, only prefix, we need to check equality here and not just prefix
		if attrs.Name == bucket {
			return nil
		}
	}

}

func (cb *GoogleCloudBuilder) createBucketIfNotExists(ctx context.Context, bucket string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return errors.Wrap(err, "getting storage client")
	}
	defer c.Close()

	_, err = c.Bucket(bucket).Attrs(ctx)

	if err == nil {
		// Bucket exists
		return nil
	}

	if err != cstorage.ErrBucketNotExist {
		return errors.Wrapf(err, "getting bucket %s", bucket)
	}

	if err := c.Bucket(bucket).Create(ctx, cb.GoogleCloudBuild.ProjectID, &cstorage.BucketAttrs{
		Name: bucket,
	}); err != nil {
		return err
	}
	logrus.Debugf("Created bucket %s in %s", bucket, cb.GoogleCloudBuild.ProjectID)
	return nil
}
<commit_msg>Pass build args to container builder<commit_after>/*
Copyright 2018 Google LLC

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 build

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"path/filepath"
	"time"

	cstorage "cloud.google.com/go/storage"
	"golang.org/x/oauth2/google"
	cloudbuild "google.golang.org/api/cloudbuild/v1"
	"google.golang.org/api/googleapi"
	"google.golang.org/api/iterator"

	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/build/tag"
	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/config"
	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/constants"
	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/docker"
	"github.com/GoogleCloudPlatform/skaffold/pkg/skaffold/util"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

const (
	// StatusUnknown "STATUS_UNKNOWN" - Status of the build is unknown.
	StatusUnknown = "STATUS_UNKNOWN"

	// StatusQueued "QUEUED" - Build is queued; work has not yet begun.
	StatusQueued = "QUEUED"

	// StatusWorking "WORKING" - Build is being executed.
	StatusWorking = "WORKING"

	// StatusSuccess  "SUCCESS" - Build finished successfully.
	StatusSuccess = "SUCCESS"

	// StatusFailure  "FAILURE" - Build failed to complete successfully.
	StatusFailure = "FAILURE"

	// StatusInternalError  "INTERNAL_ERROR" - Build failed due to an internal cause.
	StatusInternalError = "INTERNAL_ERROR"

	// StatusTimeout  "TIMEOUT" - Build took longer than was allowed.
	StatusTimeout = "TIMEOUT"

	// StatusCancelled  "CANCELLED" - Build was canceled by a user.
	StatusCancelled = "CANCELLED"

	// RetryDelay is the time to wait in between polling the status of the cloud build
	RetryDelay = 1 * time.Second
)

type GoogleCloudBuilder struct {
	*config.BuildConfig
}

func NewGoogleCloudBuilder(cfg *config.BuildConfig) (*GoogleCloudBuilder, error) {
	return &GoogleCloudBuilder{cfg}, nil
}

func (cb *GoogleCloudBuilder) Build(ctx context.Context, out io.Writer, tagger tag.Tagger, artifacts []*config.Artifact) (*BuildResult, error) {
	client, err := google.DefaultClient(ctx, cloudbuild.CloudPlatformScope)
	if err != nil {
		return nil, errors.Wrap(err, "getting google client")
	}
	cbclient, err := cloudbuild.New(client)
	if err != nil {
		return nil, errors.Wrap(err, "getting builder")
	}
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting cloud storage client")
	}
	defer c.Close()
	builds := []Build{}
	for _, artifact := range artifacts {
		build, err := cb.buildArtifact(ctx, out, cbclient, c, artifact)
		if err != nil {
			return nil, errors.Wrapf(err, "building artifact %s", artifact.ImageName)
		}
		builds = append(builds, *build)
	}

	return &BuildResult{
		Builds: builds,
	}, nil
}

func (cb *GoogleCloudBuilder) buildArtifact(ctx context.Context, out io.Writer, cbclient *cloudbuild.Service, c *cstorage.Client, artifact *config.Artifact) (*Build, error) {
	if artifact.DockerArtifact == nil {
		artifact.DockerArtifact = config.DefaultDockerArtifact
	}
	logrus.Infof("Building artifact: %+v", artifact)

	// need to format build args as strings to pass to container builder docker
	var buildArgs []string
	for k, v := range artifact.DockerArtifact.BuildArgs {
		if v != nil {
			buildArgs = append(buildArgs, []string{"--build-arg", fmt.Sprintf("%s=%s", k, *v)}...)
		}
	}
	logrus.Debugf("Build args: %s", buildArgs)

	cbBucket := fmt.Sprintf("%s%s", cb.GoogleCloudBuild.ProjectID, constants.GCSBucketSuffix)
	buildObject := fmt.Sprintf("source/%s-%s.tar.gz", cb.GoogleCloudBuild.ProjectID, util.RandomID())

	if err := cb.createBucketIfNotExists(ctx, cbBucket); err != nil {
		return nil, errors.Wrap(err, "creating bucket if not exists")
	}
	if err := cb.checkBucketProjectCorrect(ctx, cbBucket); err != nil {
		return nil, errors.Wrap(err, "checking bucket is in correct project")
	}

	fmt.Fprintf(out, "Pushing code to gs://%s/%s\n", cbBucket, buildObject)
	if err := cb.uploadTarToGCS(ctx, artifact.DockerArtifact.DockerfilePath, artifact.Workspace, cbBucket, buildObject); err != nil {
		return nil, errors.Wrap(err, "uploading source tarball")
	}

	args := append([]string{"build", "--tag", artifact.ImageName, "-f", artifact.DockerArtifact.DockerfilePath}, buildArgs...)
	args = append(args, ".")
	call := cbclient.Projects.Builds.Create(cb.GoogleCloudBuild.ProjectID, &cloudbuild.Build{
		LogsBucket: cbBucket,
		Source: &cloudbuild.Source{
			StorageSource: &cloudbuild.StorageSource{
				Bucket: cbBucket,
				Object: buildObject,
			},
		},
		Steps: []*cloudbuild.BuildStep{
			{
				Name: "gcr.io/cloud-builders/docker",
				Args: args,
			},
		},
		Images: []string{artifact.ImageName},
	})
	op, err := call.Context(ctx).Do()
	if err != nil {
		return nil, errors.Wrap(err, "could not create build")
	}

	remoteID, err := getBuildID(op)
	if err != nil {
		return nil, errors.Wrapf(err, "getting build ID from op")
	}
	logsObject := fmt.Sprintf("log-%s.txt", remoteID)
	fmt.Fprintf(out, "Logs at available at \nhttps://console.cloud.google.com/m/cloudstorage/b/%s/o/%s\n", cbBucket, logsObject)
	var imageID string
	offset := int64(0)
watch:
	for {
		logrus.Debugf("current offset %d", offset)
		b, err := cbclient.Projects.Builds.Get(cb.GoogleCloudBuild.ProjectID, remoteID).Do()
		if err != nil {
			return nil, errors.Wrap(err, "getting build status")
		}

		r, err := cb.getLogs(ctx, offset, cbBucket, logsObject)
		if err != nil {
			return nil, errors.Wrap(err, "getting logs")
		}
		if r != nil {
			written, err := io.Copy(out, r)
			if err != nil {
				return nil, errors.Wrap(err, "copying logs to stdout")
			}
			offset += written
			r.Close()
		}
		switch b.Status {
		case StatusQueued, StatusWorking, StatusUnknown:
		case StatusSuccess:
			imageID, err = getImageID(b)
			if err != nil {
				return nil, errors.Wrap(err, "getting image id from finished build")
			}
			break watch
		case StatusFailure, StatusInternalError, StatusTimeout, StatusCancelled:
			return nil, fmt.Errorf("cloud build failed: %s", b.Status)
		default:
			return nil, fmt.Errorf("unknown status: %s", b.Status)
		}

		time.Sleep(RetryDelay)
	}

	if err := c.Bucket(cbBucket).Object(buildObject).Delete(ctx); err != nil {
		return nil, errors.Wrap(err, "cleaning up source tar after build")
	}
	logrus.Infof("Deleted object %s", buildObject)
	tag := fmt.Sprintf("%s@%s", artifact.ImageName, imageID)
	logrus.Infof("Image built at %s", tag)
	return &Build{
		ImageName: artifact.ImageName,
		Tag:       tag,
		Artifact:  artifact,
	}, nil
}

func getBuildID(op *cloudbuild.Operation) (string, error) {
	if op.Metadata == nil {
		return "", errors.New("missing Metadata in operation")
	}
	var buildMeta cloudbuild.BuildOperationMetadata
	if err := json.Unmarshal([]byte(op.Metadata), &buildMeta); err != nil {
		return "", err
	}
	if buildMeta.Build == nil {
		return "", errors.New("missing Build in operation metadata")
	}
	return buildMeta.Build.Id, nil
}

func getImageID(b *cloudbuild.Build) (string, error) {
	if b.Results == nil || len(b.Results.Images) == 0 {
		return "", errors.New("build failed")
	}
	return b.Results.Images[0].Digest, nil
}

func (cb *GoogleCloudBuilder) uploadTarToGCS(ctx context.Context, dockerfilePath, dockerCtx, bucket, objectName string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return err
	}
	defer c.Close()

	relDockerfilePath := filepath.Join(dockerCtx, dockerfilePath)
	w := c.Bucket(bucket).Object(objectName).NewWriter(ctx)
	if err := docker.CreateDockerTarGzContext(w, relDockerfilePath, dockerCtx); err != nil {
		return errors.Wrap(err, "uploading targz to google storage")
	}
	return w.Close()
}

func (cb *GoogleCloudBuilder) getLogs(ctx context.Context, offset int64, bucket, objectName string) (io.ReadCloser, error) {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "getting storage client")
	}
	defer c.Close()

	r, err := c.Bucket(bucket).Object(objectName).NewRangeReader(ctx, offset, -1)
	if err != nil {
		if gerr, ok := err.(*googleapi.Error); ok {
			switch gerr.Code {
			case 404, 416, 429, 503:
				logrus.Debugf("Status Code: %d, %s", gerr.Code, gerr.Body)
				return nil, nil
			}
		}
		if err == cstorage.ErrObjectNotExist {
			logrus.Debugf("Logs for %s %s not uploaded yet...", bucket, objectName)
			return nil, nil
		}
		return nil, errors.Wrap(err, "unknown error")
	}
	return r, nil
}

func (cb *GoogleCloudBuilder) checkBucketProjectCorrect(ctx context.Context, bucket string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return errors.Wrap(err, "getting storage client")
	}
	it := c.Buckets(ctx, cb.GoogleCloudBuild.ProjectID)
	// Set the prefix to the bucket we're looking for to only return that bucket and buckets with that prefix
	// that we'll filter further later on
	it.Prefix = bucket
	for {
		attrs, err := it.Next()
		if err == iterator.Done {
			return errors.Wrap(err, "bucket not found")
		}
		if err != nil {
			return errors.Wrap(err, "iterating over buckets")
		}
		// Since we can't filter on bucket name specifically, only prefix, we need to check equality here and not just prefix
		if attrs.Name == bucket {
			return nil
		}
	}

}

func (cb *GoogleCloudBuilder) createBucketIfNotExists(ctx context.Context, bucket string) error {
	c, err := cstorage.NewClient(ctx)
	if err != nil {
		return errors.Wrap(err, "getting storage client")
	}
	defer c.Close()

	_, err = c.Bucket(bucket).Attrs(ctx)

	if err == nil {
		// Bucket exists
		return nil
	}

	if err != cstorage.ErrBucketNotExist {
		return errors.Wrapf(err, "getting bucket %s", bucket)
	}

	if err := c.Bucket(bucket).Create(ctx, cb.GoogleCloudBuild.ProjectID, &cstorage.BucketAttrs{
		Name: bucket,
	}); err != nil {
		return err
	}
	logrus.Debugf("Created bucket %s in %s", bucket, cb.GoogleCloudBuild.ProjectID)
	return nil
}
<|endoftext|>


<commit_before>package redisc

import (
	"errors"
	"strconv"
	"sync"

	"github.com/garyburd/redigo/redis"
)

const hashSlots = 16384

// Cluster manages a redis cluster. If the CreatePool field is not nil,
// a redis.Pool is used for each node in the cluster to get connections
// via Cluster.Get. If it is nil or if Cluster.Dial is called, redis.Dial
// is used to get the connection.
type Cluster struct {
	// StartupNodes is the list of initial nodes that make up
	// the cluster. The values are expected as "address:port"
	// (e.g.: "111.222.333.444:6379").
	StartupNodes []string

	// DialOptions is the list of options to set on each new connection.
	DialOptions []redis.DialOption

	// CreatePool is the function to call to create a redis.Pool for
	// the specified TCP address, using the provided options
	// as set in DialOptions. If this field is not nil, a
	// redis.Pool is created for each node in the cluster and the
	// pool is used to manage the connections unless Cluster.Dial
	// is called.
	CreatePool func(address string, options ...redis.DialOption) (*redis.Pool, error)

	// MaxAttempts is the maximum number of attempts allowed when
	// running a command. If the cluster is moving slots around in
	// its nodes, it may reply to a command with a MOVED or ASK error,
	// in which case the package retries with the redis-specified
	// node. This field controls how many of those attempts are executed
	// before returning an error.
	MaxAttempts int

	mu            sync.Mutex             // protects following fields
	err           error                  // broken connection error
	pools         map[string]*redis.Pool // created pools per node
	nodes         map[string]bool        // set of known active nodes, kept up-to-date
	mapping       [hashSlots]string      // hash slot number to master server address
	refreshNeeded bool                   // refresh mapping on next command
}

// Refresh updates the cluster's internal mapping of hash slots
// to redis node. It calls CLUSTER SLOTS on each known node until one
// of them succeeds.
//
// It should typically be called after creating the Cluster and before
// using it. The cluster automatically keeps its mapping up-to-date
// afterwards, based on the redis commands' MOVED responses.
func (c *Cluster) Refresh() error {
	c.mu.Lock()
	if err := c.err; err != nil {
		c.mu.Unlock()
		return err
	}

	// populate nodes lazily, only once
	if c.nodes == nil {
		c.populateNodes()
	}

	// grab a slice of addresses so we don't hold on to the lock during
	// the CLUSTER SLOTS calls.
	addrs := make([]string, 0, len(c.nodes))
	for addr := range c.nodes {
		addrs = append(addrs, addr)
	}
	c.mu.Unlock()

	return c.refresh(addrs)
}

func (c *Cluster) refresh(addrs []string) error {
	for _, addr := range addrs {
		m, err := c.getClusterSlots(addr)
		if err == nil {
			// succeeded, save as mapping
			c.mu.Lock()
			// mark all current nodes as false
			for k := range c.nodes {
				c.nodes[k] = false
			}
			for _, sm := range m {
				for ix := sm.start; ix <= sm.end; ix++ {
					c.mapping[ix] = sm.master
					c.nodes[sm.master] = true
				}
			}
			// remove all nodes that are gone from the cluster
			for k, ok := range c.nodes {
				if !ok {
					delete(c.nodes, k)
				}
			}
			c.mu.Unlock()

			return nil
		}
	}
	return errors.New("redisc: all nodes failed")
}

func (c *Cluster) getConnForAddr(addr string) (redis.Conn, error) {
	// non-pooled doesn't require a lock
	if c.CreatePool == nil {
		return redis.Dial("tcp", addr, c.DialOptions...)
	}

	c.mu.Lock()

	p := c.pools[addr]
	if p == nil {
		c.mu.Unlock()
		pool, err := c.CreatePool(addr, c.DialOptions...)
		if err != nil {
			return nil, err
		}

		c.mu.Lock()
		c.pools[addr] = pool
		p = pool
	}

	c.mu.Unlock()

	conn := p.Get()
	return conn, conn.Err()
}

type slotMapping struct {
	start, end int
	master     string
}

func (c *Cluster) getClusterSlots(addr string) ([]slotMapping, error) {
	conn, err := c.getConnForAddr(addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	vals, err := redis.Values(conn.Do("CLUSTER", "SLOTS"))
	if err != nil {
		return nil, err
	}

	m := make([]slotMapping, 0, len(vals))
	for len(vals) > 0 {
		var slotRange []interface{}
		vals, err = redis.Scan(vals, &slotRange)
		if err != nil {
			return nil, err
		}

		var start, end int
		var node []interface{}
		if _, err = redis.Scan(slotRange, &start, &end, &node); err != nil {
			return nil, err
		}
		var addr string
		var port int
		if _, err = redis.Scan(node, &addr, &port); err != nil {
			return nil, err
		}

		m = append(m, slotMapping{start: start, end: end, master: addr + ":" + strconv.Itoa(port)})
	}

	return m, nil
}

func (c *Cluster) populateNodes() {
	c.nodes = make(map[string]bool)
	for _, n := range c.StartupNodes {
		c.nodes[n] = true
	}
}

// Dial returns a connection the same way as Cluster.Get, but
// it guarantees that the connection will not be managed by the
// pool, even if Cluster.CreatePool is set. The actual returned
// type is *redisc.Conn, see its documentation for details.
func (c *Cluster) Dial() (redis.Conn, error) {
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()

	if err != nil {
		return nil, err
	}

	return &Conn{
		cluster:   c,
		forceDial: true,
	}, nil
}

// Get returns a redis.Conn interface that can be used to call
// redis commands on the cluster. The application must close the
// returned connection. The actual returned type is *redisc.Conn,
// see its documentation for details.
func (c *Cluster) Get() redis.Conn {
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()

	return &Conn{
		cluster: c,
		err:     err,
	}
}

// Close releases the resources used by the cluster. It closes all the
// pools that were created, if any.
func (c *Cluster) Close() error {
	c.mu.Lock()
	err := c.err
	if err == nil {
		c.err = errors.New("redisc: closed")
		for _, p := range c.pools {
			if e := p.Close(); e != nil && err == nil {
				err = e
			}
		}
	}
	c.mu.Unlock()

	return err
}
<commit_msg>redisc: handle getting a connection<commit_after>package redisc

import (
	"errors"
	"math/rand"
	"strconv"
	"sync"
	"time"

	"github.com/garyburd/redigo/redis"
)

var rnd = rand.New(rand.NewSource(time.Now().UnixNano()))

const hashSlots = 16384

// Cluster manages a redis cluster. If the CreatePool field is not nil,
// a redis.Pool is used for each node in the cluster to get connections
// via Cluster.Get. If it is nil or if Cluster.Dial is called, redis.Dial
// is used to get the connection.
type Cluster struct {
	// StartupNodes is the list of initial nodes that make up
	// the cluster. The values are expected as "address:port"
	// (e.g.: "111.222.333.444:6379").
	StartupNodes []string

	// DialOptions is the list of options to set on each new connection.
	DialOptions []redis.DialOption

	// CreatePool is the function to call to create a redis.Pool for
	// the specified TCP address, using the provided options
	// as set in DialOptions. If this field is not nil, a
	// redis.Pool is created for each node in the cluster and the
	// pool is used to manage the connections unless Cluster.Dial
	// is called.
	CreatePool func(address string, options ...redis.DialOption) (*redis.Pool, error)

	// MaxAttempts is the maximum number of attempts allowed when
	// running a command on a connection returned by RetryConn, in
	// order to automatically follow MOVED or ASK redirections. This
	// field controls how many of those attempts are executed before
	// returning an error.
	MaxAttempts int

	mu            sync.Mutex             // protects following fields
	err           error                  // broken connection error
	pools         map[string]*redis.Pool // created pools per node
	nodes         map[string]bool        // set of known active nodes, kept up-to-date
	mapping       [hashSlots]string      // hash slot number to master server address
	refreshNeeded bool                   // refresh mapping on next command
}

// Refresh updates the cluster's internal mapping of hash slots
// to redis node. It calls CLUSTER SLOTS on each known node until one
// of them succeeds.
//
// It should typically be called after creating the Cluster and before
// using it. The cluster automatically keeps its mapping up-to-date
// afterwards, based on the redis commands' MOVED responses.
func (c *Cluster) Refresh() error {
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()
	if err != nil {
		return err
	}

	return c.refresh()
}

func (c *Cluster) refresh() error {
	addrs := c.getNodeAddrs()
	for _, addr := range addrs {
		m, err := c.getClusterSlots(addr)
		if err == nil {
			// succeeded, save as mapping
			c.mu.Lock()
			// mark all current nodes as false
			for k := range c.nodes {
				c.nodes[k] = false
			}
			for _, sm := range m {
				for ix := sm.start; ix <= sm.end; ix++ {
					c.mapping[ix] = sm.master
					c.nodes[sm.master] = true
				}
			}
			// remove all nodes that are gone from the cluster
			for k, ok := range c.nodes {
				if !ok {
					delete(c.nodes, k)
				}
			}
			// mark that no refresh is needed until another MOVED
			c.refreshNeeded = false
			c.mu.Unlock()

			return nil
		}
	}
	return errors.New("redisc: all nodes failed")
}

type slotMapping struct {
	start, end int
	master     string
}

func (c *Cluster) getClusterSlots(addr string) ([]slotMapping, error) {
	conn, err := c.getConnForAddr(addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	vals, err := redis.Values(conn.Do("CLUSTER", "SLOTS"))
	if err != nil {
		return nil, err
	}

	m := make([]slotMapping, 0, len(vals))
	for len(vals) > 0 {
		var slotRange []interface{}
		vals, err = redis.Scan(vals, &slotRange)
		if err != nil {
			return nil, err
		}

		var start, end int
		var node []interface{}
		if _, err = redis.Scan(slotRange, &start, &end, &node); err != nil {
			return nil, err
		}
		var addr string
		var port int
		if _, err = redis.Scan(node, &addr, &port); err != nil {
			return nil, err
		}

		m = append(m, slotMapping{start: start, end: end, master: addr + ":" + strconv.Itoa(port)})
	}

	return m, nil
}

func (c *Cluster) getConnForAddr(addr string) (redis.Conn, error) {
	// non-pooled doesn't require a lock
	if c.CreatePool == nil {
		return redis.Dial("tcp", addr, c.DialOptions...)
	}

	c.mu.Lock()

	p := c.pools[addr]
	if p == nil {
		c.mu.Unlock()
		pool, err := c.CreatePool(addr, c.DialOptions...)
		if err != nil {
			return nil, err
		}

		c.mu.Lock()
		c.pools[addr] = pool
		p = pool
	}

	c.mu.Unlock()

	conn := p.Get()
	return conn, conn.Err()
}

func (c *Cluster) getConnForSlot(slot int) (redis.Conn, error) {
	c.mu.Lock()
	addr := c.mapping[slot]
	c.mu.Unlock()
	if addr == "" {
		return nil, errors.New("redisc: no node for slot " + strconv.Itoa(slot))
	}
	return c.getConnForAddr(addr)
}

func (c *Cluster) getRandomConn() (redis.Conn, error) {
	addrs := c.getNodeAddrs()
	perms := rnd.Perm(len(addrs))
	for _, ix := range perms {
		addr := addrs[ix]
		conn, err := c.getConnForAddr(addr)
		if err == nil {
			return conn, nil
		}
	}
	return nil, errors.New("redisc: failed to get a connection")
}

func (c *Cluster) getConn(preferredSlot int) (redis.Conn, error) {
	conn, err := c.getConnForSlot(preferredSlot)
	if err != nil {
		conn, err = c.getRandomConn()
	}
	return conn, err
}

func (c *Cluster) getNodeAddrs() []string {
	c.mu.Lock()

	// populate nodes lazily, only once
	if c.nodes == nil {
		c.nodes = make(map[string]bool)
		for _, n := range c.StartupNodes {
			c.nodes[n] = true
		}
	}

	// grab a slice of addresses
	addrs := make([]string, 0, len(c.nodes))
	for addr := range c.nodes {
		addrs = append(addrs, addr)
	}
	c.mu.Unlock()

	return addrs
}

// Dial returns a connection the same way as Cluster.Get, but
// it guarantees that the connection will not be managed by the
// pool, even if Cluster.CreatePool is set. The actual returned
// type is *redisc.Conn, see its documentation for details.
func (c *Cluster) Dial() (redis.Conn, error) {
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()

	if err != nil {
		return nil, err
	}

	return &Conn{
		cluster:   c,
		forceDial: true,
	}, nil
}

// Get returns a redis.Conn interface that can be used to call
// redis commands on the cluster. The application must close the
// returned connection. The actual returned type is *redisc.Conn,
// see its documentation for details.
func (c *Cluster) Get() redis.Conn {
	c.mu.Lock()
	err := c.err
	c.mu.Unlock()

	return &Conn{
		cluster: c,
		err:     err,
	}
}

// Close releases the resources used by the cluster. It closes all the
// pools that were created, if any.
func (c *Cluster) Close() error {
	c.mu.Lock()
	err := c.err
	if err == nil {
		c.err = errors.New("redisc: closed")
		for _, p := range c.pools {
			if e := p.Close(); e != nil && err == nil {
				err = e
			}
		}
	}
	c.mu.Unlock()

	return err
}
<|endoftext|>