<commit_before>package emailnotifier

import (
	"bytes"
	"fmt"
	"html/template"
	"net/url"
	"socialapi/config"
	"socialapi/workers/notification/models"
)

var (
	mainTemplateFile        = "../socialapi/workers/emailnotifier/templates/main.tmpl"
	footerTemplateFile      = "../socialapi/workers/emailnotifier/templates/footer.tmpl"
	contentTemplateFile     = "../socialapi/workers/emailnotifier/templates/content.tmpl"
	gravatarTemplateFile    = "../socialapi/workers/emailnotifier/templates/gravatar.tmpl"
	groupTemplateFile       = "../socialapi/workers/emailnotifier/templates/group.tmpl"
	previewTemplateFile     = "../socialapi/workers/emailnotifier/templates/preview.tmpl"
	objectTemplateFile      = "../socialapi/workers/emailnotifier/templates/object.tmpl"
	unsubscribeTemplateFile = "../socialapi/workers/emailnotifier/templates/unsubscribe.tmpl"
)

type NotificationContainer struct {
	Activity        *models.NotificationActivity
	Content         *models.NotificationContent
	Notification    *models.Notification
	Message         string
	Slug            string
	ActivityMessage string
	ObjectType      string
	Group           GroupContent
	Token           string
}

type MailContent struct {
	CurrentDate  string
	FirstName    string
	Description  string
	Size         int
	Uri          string
	ActorContact UserContact
	// event content
	ActivityTime string
	Avatar       string
	Sender       string
	Action       string
	ContentLink  string
	Group        GroupContent
	Message      string
	Slug         string
	ObjectType   string
	Token        string
	ContentType  string
	Recipient    string
}

type GroupContent struct {
	Slug string
	Name string
}

func renderTemplate(uc *UserContact, nc *NotificationContainer) (string, error) {
	t := template.Must(template.ParseFiles(
		mainTemplateFile, footerTemplateFile, contentTemplateFile,
		gravatarTemplateFile, unsubscribeTemplateFile))
	mc, err := buildMailContent(uc, nc)
	t = appendPreviewTemplate(t, nc)
	t = appendGroupTemplate(t, nc)

	if err != nil {
		return "", err
	}

	var doc bytes.Buffer
	if err := t.Execute(&doc, mc); err != nil {
		return "", err
	}

	return doc.String(), nil
}

func buildMailContent(uc *UserContact, nc *NotificationContainer) (*MailContent, error) {
	mc := &MailContent{
		FirstName:   uc.FirstName,
		Size:        20,
		Slug:        nc.Slug,
		Message:     nc.Message,
		Action:      nc.ActivityMessage,
		ObjectType:  nc.ObjectType,
		Group:       nc.Group,
		Token:       nc.Token,
		ContentType: nc.Content.TypeConstant,
		Recipient:   url.QueryEscape(uc.Email),
	}

	mc.CurrentDate, mc.ActivityTime = prepareDateTime(nc)
	mc.Uri = config.Get().Uri

	actor, err := fetchUserContact(nc.Activity.ActorId)
	if err != nil {
		return nil, fmt.Errorf("an error occurred while retrieving actor details", err)
	}
	mc.ActorContact = *actor

	return mc, nil
}

func appendGroupTemplate(t *template.Template, nc *NotificationContainer) *template.Template {
	var groupTemplate *template.Template
	if nc.Group.Name == "" || nc.Group.Slug == "koding" {
		groupTemplate = getEmptyTemplate()
	} else {
		groupTemplate = template.Must(
			template.ParseFiles(groupTemplateFile))
	}

	t.AddParseTree("group", groupTemplate.Tree)

	return t
}

func appendPreviewTemplate(t *template.Template, nc *NotificationContainer) *template.Template {
	var previewTemplate, objectTemplate *template.Template
	if nc.Message == "" {
		previewTemplate = getEmptyTemplate()
		objectTemplate = getEmptyTemplate()
	} else {
		previewTemplate = template.Must(template.ParseFiles(previewTemplateFile))
		objectTemplate = template.Must(template.ParseFiles(objectTemplateFile))
	}

	t.AddParseTree("preview", previewTemplate.Tree)
	t.AddParseTree("object", objectTemplate.Tree)

	return t
}

func getEmptyTemplate() *template.Template {
	return template.Must(template.New("").Parse(""))
}

func prepareDateTime(nc *NotificationContainer) (string, string) {
	layoutDate := "Jan 02"
	layoutTime := "15:04"

	return nc.Notification.ActivatedAt.Format(layoutDate),
		nc.Notification.ActivatedAt.Format(layoutTime)
}

func prepareSubject(nc *NotificationContainer) string {
	t, err := nc.Content.GetContentType()
	if err != nil {
		return ""
	}

	return t.GetDefinition()
}
<commit_msg>Notification: MailContent is breakdown into three structs as EventContent, UnsubscribeContent and MailContent<commit_after>package emailnotifier

import (
	"bytes"
	"fmt"
	"html/template"
	"net/url"
	"socialapi/config"
	"socialapi/workers/notification/models"
	"time"
)

var (
	mainTemplateFile        = "../socialapi/workers/emailnotifier/templates/main.tmpl"
	footerTemplateFile      = "../socialapi/workers/emailnotifier/templates/footer.tmpl"
	contentTemplateFile     = "../socialapi/workers/emailnotifier/templates/content.tmpl"
	gravatarTemplateFile    = "../socialapi/workers/emailnotifier/templates/gravatar.tmpl"
	groupTemplateFile       = "../socialapi/workers/emailnotifier/templates/group.tmpl"
	previewTemplateFile     = "../socialapi/workers/emailnotifier/templates/preview.tmpl"
	objectTemplateFile      = "../socialapi/workers/emailnotifier/templates/object.tmpl"
	unsubscribeTemplateFile = "../socialapi/workers/emailnotifier/templates/unsubscribe.tmpl"
)

type NotificationContainer struct {
	Activity        *models.NotificationActivity
	Content         *models.NotificationContent
	AccountId       int64
	Message         string
	Slug            string
	ActivityMessage string
	ObjectType      string
	Group           GroupContent
	CreatedAt       time.Time
}

type EventContent struct {
	// event content
	ActivityTime string
	ActorContact UserContact
	Action       string
	Size         int
	Slug         string
	Uri          string
	ObjectType   string
	Group        GroupContent
	Message      string
}

type MailContent struct {
	CurrentDate string
	FirstName   string
	Description string
	Uri         string
	ContentLink string
	Content     template.HTML
	Unsubscribe *UnsubscribeContent
}

type UnsubscribeContent struct {
	Token       string
	ContentType string
	Recipient   string
}

type GroupContent struct {
	Slug string
	Name string
}

func renderTemplate(uc *UserContact, nc *NotificationContainer) (string, error) {
	t := template.Must(template.ParseFiles(
		mainTemplateFile, footerTemplateFile, contentTemplateFile,
		gravatarTemplateFile, unsubscribeTemplateFile))
	mc, err := buildMailContent(uc, nc)
	t = appendPreviewTemplate(t, nc)
	t = appendGroupTemplate(t, nc)

	if err != nil {
		return "", err
	}

	var doc bytes.Buffer
	if err := t.Execute(&doc, mc); err != nil {
		return "", err
	}

	return doc.String(), nil
}

func buildMailContent(uc *UserContact, nc *NotificationContainer) (*MailContent, error) {
	mc := &MailContent{
		FirstName:   uc.FirstName,
		Size:        20,
		Slug:        nc.Slug,
		Message:     nc.Message,
		Action:      nc.ActivityMessage,
		ObjectType:  nc.ObjectType,
		Group:       nc.Group,
		Token:       nc.Token,
		ContentType: nc.Content.TypeConstant,
		Recipient:   url.QueryEscape(uc.Email),
	}

	mc.CurrentDate, mc.ActivityTime = prepareDateTime(nc)
	mc.Uri = config.Get().Uri

	actor, err := fetchUserContact(nc.Activity.ActorId)
	if err != nil {
		return nil, fmt.Errorf("an error occurred while retrieving actor details", err)
	}
	mc.ActorContact = *actor

	return mc, nil
}

func appendGroupTemplate(t *template.Template, nc *NotificationContainer) *template.Template {
	var groupTemplate *template.Template
	if nc.Group.Name == "" || nc.Group.Slug == "koding" {
		groupTemplate = getEmptyTemplate()
	} else {
		groupTemplate = template.Must(
			template.ParseFiles(groupTemplateFile))
	}

	t.AddParseTree("group", groupTemplate.Tree)

	return t
}

func appendPreviewTemplate(t *template.Template, nc *NotificationContainer) *template.Template {
	var previewTemplate, objectTemplate *template.Template
	if nc.Message == "" {
		previewTemplate = getEmptyTemplate()
		objectTemplate = getEmptyTemplate()
	} else {
		previewTemplate = template.Must(template.ParseFiles(previewTemplateFile))
		objectTemplate = template.Must(template.ParseFiles(objectTemplateFile))
	}

	t.AddParseTree("preview", previewTemplate.Tree)
	t.AddParseTree("object", objectTemplate.Tree)

	return t
}

func getEmptyTemplate() *template.Template {
	return template.Must(template.New("").Parse(""))
}

func prepareDateTime(nc *NotificationContainer) (string, string) {
	layoutDate := "Jan 02"
	layoutTime := "15:04"

	return nc.Notification.ActivatedAt.Format(layoutDate),
		nc.Notification.ActivatedAt.Format(layoutTime)
}

func prepareSubject(nc *NotificationContainer) string {
	t, err := nc.Content.GetContentType()
	if err != nil {
		return ""
	}

	return t.GetDefinition()
}
<|endoftext|>


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

package fuse

import (
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
	"time"
)

func TestFuse(t *testing.T) {
	dir := "/mnt"
	exec.Command("umount", dir).Run()
	os.MkdirAll(dir, 0777)

	c, err := Mount(dir)
	if err != nil {
		t.Fatal(err)
	}
	defer exec.Command("umount", dir).Run()

	go c.Serve(testFS{})
	time.Sleep(1*time.Second)

	_, err = os.Stat(dir+"/"+fuseTests[0].name)
	if err != nil {
		t.Fatalf("mount did not work")
		return
	}

	for _, tt := range fuseTests {
		tt.node.test(dir+"/"+tt.name, t)
	}
}

var fuseTests = []struct{
	name string
	node interface { Node; test(string, *testing.T) }
}{
	{"readAll", readAll{}},
	{"writeAll", &writeAll{}},
}

type readAll struct{ file }

const hi = "hello, world"
func (readAll) ReadAll(intr Intr) ([]byte, Error) {
	return []byte(hi), nil
}

func (readAll) test(path string, t *testing.T) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		t.Errorf("readAll: %v", err)
		return
	}
	if string(data) != hi {
		t.Errorf("readAll = %q, want %q", data, hi)
	}
}

type writeAll struct {
	file
	data []byte
}

func (w *writeAll) WriteAll(data []byte, intr Intr) Error {
	w.data = data
	return nil
}

func (w *writeAll) test(path string, t *testing.T) {
	err := ioutil.WriteFile(path, []byte(hi), 0666)
	if err != nil {
		t.Errorf("WriteFile: %v", err)
		return
	}
	if string(w.data) != hi {
		t.Errorf("writeAll = %q, want %q", w.data, hi)
	}
}

type file struct {}
type dir struct {}

func (f file) Attr() Attr { return Attr{Mode: 0666} }
func (f dir) Attr() Attr { return Attr{Mode: os.ModeDir | 0777} }

type testFS struct {}

func (testFS) Root() (Node, Error) {
	return testFS{}, nil
}

func (testFS) Attr() Attr {
	return Attr{Mode: os.ModeDir | 0555}
}

func (testFS) Lookup(name string, intr Intr) (Node, Error) {
	for _, tt := range fuseTests {
		if tt.name == name {
			return tt.node, nil
		}
	}
	return nil, ENOENT
}

func (testFS) ReadDir(intr Intr) ([]Dirent, Error) {
	var dirs []Dirent
	for _, tt := range fuseTests {
		dirs = append(dirs, Dirent{Name: tt.name})
	}
	return dirs, nil
}
<commit_msg>fuse: more tests<commit_after>// Copyright 2012 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package fuse

import (
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
	"time"
)

func TestFuse(t *testing.T) {
	dir := "/tmp/fusetestmnt"
	exec.Command("umount", dir).Run()
	os.MkdirAll(dir, 0777)

	c, err := Mount(dir)
	if err != nil {
		t.Fatal(err)
	}
	defer exec.Command("umount", dir).Run()

	go c.Serve(testFS{})
	time.Sleep(1*time.Second)

	_, err = os.Stat(dir+"/"+fuseTests[0].name)
	if err != nil {
		t.Fatalf("mount did not work")
		return
	}

	for _, tt := range fuseTests {
		tt.node.test(dir+"/"+tt.name, t)
	}
}

var fuseTests = []struct{
	name string
	node interface { Node; test(string, *testing.T) }
}{
	{"readAll", readAll{}},
	{"readAll1", &readAll1{}},
	{"writeAll", &writeAll{}},
	{"writeAll2", &writeAll2{}},
	{"release", &release{}},
}

// TO TEST:
//	Statfs
//	Lookup(*LookupRequest, *LookupResponse)
//	Getattr(*GetattrRequest, *GetattrResponse)
//	Attr with explicit inode
//	Setattr(*SetattrRequest, *SetattrResponse)
//	Access(*AccessRequest)
//	Mkdir(*MkdirRequest)
//	Open(*OpenRequest, *OpenResponse)
//	Create(*CreateRequest, *CreateResponse)
//	Getxattr, Setxattr, Listxattr, Removexattr
//	Write(*WriteRequest, *WriteResponse)
//	Flush(*FlushRequest, *FlushResponse)

// Test Read calling ReadAll.

type readAll struct{ file }

const hi = "hello, world"
func (readAll) ReadAll(intr Intr) ([]byte, Error) {
	return []byte(hi), nil
}

func (readAll) test(path string, t *testing.T) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		t.Errorf("readAll: %v", err)
		return
	}
	if string(data) != hi {
		t.Errorf("readAll = %q, want %q", data, hi)
	}
}

// Test Read.

type readAll1 struct { file }

func (readAll1) Read(req *ReadRequest, resp *ReadResponse, intr Intr) Error {
	HandleRead(req, resp, []byte(hi))
	return nil
}

func (readAll1) test(path string, t *testing.T) {
	readAll{}.test(path, t)
}

// Test Write calling WriteAll.

type writeAll struct {
	file
	data []byte
}

func (w *writeAll) WriteAll(data []byte, intr Intr) Error {
	w.data = data
	return nil
}

func (w *writeAll) test(path string, t *testing.T) {
	err := ioutil.WriteFile(path, []byte(hi), 0666)
	if err != nil {
		t.Errorf("WriteFile: %v", err)
		return
	}
	if string(w.data) != hi {
		t.Errorf("writeAll = %q, want %q", w.data, hi)
	}
}

// Test Write calling Setattr+Write+Flush.

type writeAll2 struct {
	file
	data []byte
	setattr bool
	flush bool
}

func (w *writeAll2) Setattr(req *SetattrRequest, resp *SetattrResponse, intr Intr) Error {
	w.setattr = true
	return nil
}

func (w *writeAll2) Flush(req *FlushRequest, intr Intr) Error {
	w.flush = true
	return nil
}

func (w *writeAll2) Write(req *WriteRequest, resp *WriteResponse, intr Intr) Error {
	w.data = append(w.data, req.Data...)
	resp.Size = len(req.Data)
	return nil
}

func (w *writeAll2) test(path string, t *testing.T) {
	err := ioutil.WriteFile(path, []byte(hi), 0666)
	if err != nil {
		t.Errorf("WriteFile: %v", err)
		return
	}
	if !w.setattr || string(w.data) != hi || !w.flush {
		t.Errorf("writeAll = %v, %q, %v, want %v, %q, %v", w.setattr, string(w.data), w.flush, true, hi, true)
	}
}

// Test Release.

type release struct {
	file
	did bool
}

func (r *release) Release(*ReleaseRequest, Intr) Error {
	r.did = true
	return nil
}

func (r *release) test(path string, t *testing.T) {
	r.did = false
	f, err := os.Open(path)
	if err != nil {
		t.Error(err)
		return
	}
	f.Close()
	time.Sleep(1*time.Second)
	if !r.did {
		t.Error("Close did not Release")
	}
}

type file struct {}
type dir struct {}

func (f file) Attr() Attr { return Attr{Mode: 0666} }
func (f dir) Attr() Attr { return Attr{Mode: os.ModeDir | 0777} }

type testFS struct {}

func (testFS) Root() (Node, Error) {
	return testFS{}, nil
}

func (testFS) Attr() Attr {
	return Attr{Mode: os.ModeDir | 0555}
}

func (testFS) Lookup(name string, intr Intr) (Node, Error) {
	for _, tt := range fuseTests {
		if tt.name == name {
			return tt.node, nil
		}
	}
	return nil, ENOENT
}

func (testFS) ReadDir(intr Intr) ([]Dirent, Error) {
	var dirs []Dirent
	for _, tt := range fuseTests {
		dirs = append(dirs, Dirent{Name: tt.name})
	}
	return dirs, nil
}
<|endoftext|>


<commit_before>package hatchery

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/docker/docker/pkg/namesgenerator"
	"github.com/patrickmn/go-cache"

	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

// Create creates hatchery
func Create(h Interface) {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)

	// Gracefully shutdown connections
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGKILL)
	defer func() {
		signal.Stop(c)
		cancel()
	}()

	go func() {
		select {
		case <-c:
			defer cancel()
			return
		case <-ctx.Done():
			return
		}
	}()

	if err := h.Init(); err != nil {
		log.Error("Create> Init error: %s", err)
		os.Exit(10)
	}

	hostname, errh := os.Hostname()
	if errh != nil {
		log.Error("Create> Cannot retrieve hostname: %s", errh)
		os.Exit(10)
	}

	go hearbeat(h, h.Configuration().API.Token, h.Configuration().API.MaxHeartbeatFailures)

	pbjobs := make(chan sdk.PipelineBuildJob, 1)
	wjobs := make(chan sdk.WorkflowNodeJobRun, 1)
	errs := make(chan error, 1)
	var nRoutines, workersStarted int64

	go func(ctx context.Context) {
		if err := h.Client().QueuePolling(ctx, wjobs, pbjobs, errs, 2*time.Second, h.Configuration().Provision.GraceTimeQueued); err != nil {
			log.Error("Queues polling stopped: %v", err)
		}
	}(ctx)

	// Create a cache with a default expiration time of 3 second, and which
	// purges expired items every minute
	spawnIDs := cache.New(3*time.Second, 60*time.Second)

	tickerProvision := time.NewTicker(time.Duration(h.Configuration().Provision.Frequency) * time.Second)
	tickerRegister := time.NewTicker(time.Duration(h.Configuration().Provision.RegisterFrequency) * time.Second)
	tickerCountWorkersStarted := time.NewTicker(time.Duration(2 * time.Second))
	tickerGetModels := time.NewTicker(time.Duration(3 * time.Second))

	var maxWorkersReached bool
	var models []sdk.Model

	for {
		select {
		case <-ctx.Done():
			if err := ctx.Err(); err != nil {
				log.Error("Exiting Hatchery: %v", err)
			} else {
				log.Info("Exiting Hatchery")
			}
			tickerRegister.Stop()
			return
		case <-tickerCountWorkersStarted.C:
			workersStarted = int64(h.WorkersStarted())
			if workersStarted > int64(h.Configuration().Provision.MaxWorker) {
				log.Info("max workers reached. current:%d max:%d", workersStarted, int64(h.Configuration().Provision.MaxWorker))
				maxWorkersReached = true
			} else {
				maxWorkersReached = false
			}
			log.Debug("workers already started:%d", workersStarted)
		case <-tickerGetModels.C:
			var errwm error
			models, errwm = h.Client().WorkerModelsEnabled()
			if errwm != nil {
				log.Error("error on h.Client().WorkerModelsEnabled(): %v", errwm)
			}
		case j := <-pbjobs:
			if maxWorkersReached {
				log.Debug("maxWorkerReached:%d", workersStarted)
				continue
			}
			go func(job sdk.PipelineBuildJob) {
				if isRun := receiveJob(h, false, job.ExecGroups, job.ID, job.QueuedSeconds, job.BookedBy, job.Job.Action.Requirements, models, &nRoutines, spawnIDs, hostname); isRun {
					atomic.AddInt64(&workersStarted, 1)
					spawnIDs.SetDefault(string(job.ID), job.ID)
				}
			}(j)
		case j := <-wjobs:
			if maxWorkersReached {
				log.Debug("maxWorkerReached:%d", workersStarted)
				continue
			}
			go func(job sdk.WorkflowNodeJobRun) {
				if isRun := receiveJob(h, true, nil, job.ID, job.QueuedSeconds, job.BookedBy, job.Job.Action.Requirements, models, &nRoutines, spawnIDs, hostname); isRun {
					atomic.AddInt64(&workersStarted, 1)
					spawnIDs.SetDefault(string(job.ID), job.ID)
				}
			}(j)
		case err := <-errs:
			log.Error("%v", err)
		case <-tickerProvision.C:
			provisioning(h, h.Configuration().Provision.Disabled, models)
		case <-tickerRegister.C:
			if err := workerRegister(h, models); err != nil {
				log.Warning("Error on workerRegister: %s", err)
			}
		}
	}
}

// Register calls CDS API to register current hatchery
func Register(h Interface) error {
	newHatchery, uptodate, err := h.Client().HatcheryRegister(*h.Hatchery())
	if err != nil {
		return sdk.WrapError(err, "register> Got HTTP exiting")
	}
	h.Hatchery().ID = newHatchery.ID
	h.Hatchery().GroupID = newHatchery.GroupID
	h.Hatchery().Model = newHatchery.Model
	h.Hatchery().Name = newHatchery.Name

	log.Info("Register> Hatchery %s registered with id:%d", h.Hatchery().Name, h.Hatchery().ID)

	if !uptodate {
		log.Warning("-=-=-=-=- Please update your hatchery binary -=-=-=-=-")
	}
	return nil
}

// GenerateName generate a hatchery's name
func GenerateName(add, name string) string {
	if name == "" {
		var errHostname error
		name, errHostname = os.Hostname()
		if errHostname != nil {
			log.Warning("Cannot retrieve hostname: %s", errHostname)
			name = "cds-hatchery"
		}
		name += "-" + namesgenerator.GetRandomName(0)
	}

	if add != "" {
		name += "-" + add
	}

	return name
}

func hearbeat(m Interface, token string, maxFailures int) {
	var failures int
	for {
		time.Sleep(5 * time.Second)
		if m.Hatchery().ID == 0 {
			log.Info("hearbeat> %s Disconnected from CDS engine, trying to register...", m.Hatchery().Name)
			if err := Register(m); err != nil {
				log.Info("hearbeat> %s Cannot register: %s", m.Hatchery().Name, err)
				checkFailures(maxFailures, failures)
				continue
			}
			if m.Hatchery().ID == 0 {
				log.Error("hearbeat> Cannot register hatchery. ID %d", m.Hatchery().ID)
				checkFailures(maxFailures, failures)
				continue
			}
			log.Info("hearbeat> %s Registered back: ID %d with model ID %d", m.Hatchery().Name, m.Hatchery().ID, m.Hatchery().Model.ID)
		}

		if err := m.Client().HatcheryRefresh(m.Hatchery().ID); err != nil {
			log.Info("heartbeat> %s cannot refresh beat: %s", m.Hatchery().Name, err)
			m.Hatchery().ID = 0
			checkFailures(maxFailures, failures)
			continue
		}
		failures = 0
	}
}

func checkFailures(maxFailures, nb int) {
	if nb > maxFailures {
		log.Error("Too many failures on try register. This hatchery is killed")
		os.Exit(10)
	}
}

func workerRegister(h Interface, models []sdk.Model) error {
	if len(models) == 0 {
		return fmt.Errorf("workerRegister> No model returned by GetWorkerModels")
	}
	log.Debug("workerRegister> models received: %d", len(models))

	var nRegistered int
	for _, m := range models {
		if m.Type != h.ModelType() {
			continue
		}

		// limit to 5 registration per ticker
		if nRegistered > 5 {
			break
		}
		// if current hatchery is in same group than worker model -> do not avoid spawn, even if worker model is in error
		if m.NbSpawnErr > 5 && h.Hatchery().GroupID != m.ID {
			log.Warning("workerRegister> Too many errors on spawn with model %s, please check this worker model", m.Name)
			continue
		}

		if h.NeedRegistration(&m) {
			log.Info("workerRegister> spawn a worker for register worker model %s (%d)", m.Name, m.ID)
			if _, errSpawn := h.SpawnWorker(&m, 0, nil, true, "spawn for register"); errSpawn != nil {
				log.Warning("workerRegister> cannot spawn worker for register: %s", m.Name, errSpawn)
				if err := h.Client().WorkerModelSpawnError(m.ID, fmt.Sprintf("workerRegister> cannot spawn worker for register: %s", errSpawn)); err != nil {
					log.Error("workerRegister> error on call client.WorkerModelSpawnError on worker model %s for register: %s", m.Name, errSpawn)
				}
				continue
			}
			nRegistered++
		} else {
			log.Debug("workerRegister> no need to register worker model %s (%d)", m.Name, m.ID)
			continue
		}
	}
	return nil
}
<commit_msg>fix (hatchery): avoid log at startup (#1266)<commit_after>package hatchery

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/docker/docker/pkg/namesgenerator"
	"github.com/patrickmn/go-cache"

	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

// Create creates hatchery
func Create(h Interface) {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)

	// Gracefully shutdown connections
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGKILL)
	defer func() {
		signal.Stop(c)
		cancel()
	}()

	go func() {
		select {
		case <-c:
			defer cancel()
			return
		case <-ctx.Done():
			return
		}
	}()

	if err := h.Init(); err != nil {
		log.Error("Create> Init error: %s", err)
		os.Exit(10)
	}

	hostname, errh := os.Hostname()
	if errh != nil {
		log.Error("Create> Cannot retrieve hostname: %s", errh)
		os.Exit(10)
	}

	go hearbeat(h, h.Configuration().API.Token, h.Configuration().API.MaxHeartbeatFailures)

	pbjobs := make(chan sdk.PipelineBuildJob, 1)
	wjobs := make(chan sdk.WorkflowNodeJobRun, 1)
	errs := make(chan error, 1)
	var nRoutines, workersStarted int64

	go func(ctx context.Context) {
		if err := h.Client().QueuePolling(ctx, wjobs, pbjobs, errs, 2*time.Second, h.Configuration().Provision.GraceTimeQueued); err != nil {
			log.Error("Queues polling stopped: %v", err)
		}
	}(ctx)

	// Create a cache with a default expiration time of 3 second, and which
	// purges expired items every minute
	spawnIDs := cache.New(3*time.Second, 60*time.Second)

	tickerProvision := time.NewTicker(time.Duration(h.Configuration().Provision.Frequency) * time.Second)
	tickerRegister := time.NewTicker(time.Duration(h.Configuration().Provision.RegisterFrequency) * time.Second)
	tickerCountWorkersStarted := time.NewTicker(time.Duration(2 * time.Second))
	tickerGetModels := time.NewTicker(time.Duration(3 * time.Second))

	var maxWorkersReached bool
	var models []sdk.Model

	// Call WorkerModel Enabled first
	var errwm error
	models, errwm = h.Client().WorkerModelsEnabled()
	if errwm != nil {
		log.Error("error on h.Client().WorkerModelsEnabled() (init call): %v", errwm)
	}

	for {
		select {
		case <-ctx.Done():
			if err := ctx.Err(); err != nil {
				log.Error("Exiting Hatchery: %v", err)
			} else {
				log.Info("Exiting Hatchery")
			}
			tickerRegister.Stop()
			return
		case <-tickerCountWorkersStarted.C:
			workersStarted = int64(h.WorkersStarted())
			if workersStarted > int64(h.Configuration().Provision.MaxWorker) {
				log.Info("max workers reached. current:%d max:%d", workersStarted, int64(h.Configuration().Provision.MaxWorker))
				maxWorkersReached = true
			} else {
				maxWorkersReached = false
			}
			log.Debug("workers already started:%d", workersStarted)
		case <-tickerGetModels.C:
			var errwm error
			models, errwm = h.Client().WorkerModelsEnabled()
			if errwm != nil {
				log.Error("error on h.Client().WorkerModelsEnabled(): %v", errwm)
			}
		case j := <-pbjobs:
			if maxWorkersReached {
				log.Debug("maxWorkerReached:%d", workersStarted)
				continue
			}
			go func(job sdk.PipelineBuildJob) {
				if isRun := receiveJob(h, false, job.ExecGroups, job.ID, job.QueuedSeconds, job.BookedBy, job.Job.Action.Requirements, models, &nRoutines, spawnIDs, hostname); isRun {
					atomic.AddInt64(&workersStarted, 1)
					spawnIDs.SetDefault(string(job.ID), job.ID)
				}
			}(j)
		case j := <-wjobs:
			if maxWorkersReached {
				log.Debug("maxWorkerReached:%d", workersStarted)
				continue
			}
			go func(job sdk.WorkflowNodeJobRun) {
				if isRun := receiveJob(h, true, nil, job.ID, job.QueuedSeconds, job.BookedBy, job.Job.Action.Requirements, models, &nRoutines, spawnIDs, hostname); isRun {
					atomic.AddInt64(&workersStarted, 1)
					spawnIDs.SetDefault(string(job.ID), job.ID)
				}
			}(j)
		case err := <-errs:
			log.Error("%v", err)
		case <-tickerProvision.C:
			provisioning(h, h.Configuration().Provision.Disabled, models)
		case <-tickerRegister.C:
			if err := workerRegister(h, models); err != nil {
				log.Warning("Error on workerRegister: %s", err)
			}
		}
	}
}

// Register calls CDS API to register current hatchery
func Register(h Interface) error {
	newHatchery, uptodate, err := h.Client().HatcheryRegister(*h.Hatchery())
	if err != nil {
		return sdk.WrapError(err, "register> Got HTTP exiting")
	}
	h.Hatchery().ID = newHatchery.ID
	h.Hatchery().GroupID = newHatchery.GroupID
	h.Hatchery().Model = newHatchery.Model
	h.Hatchery().Name = newHatchery.Name

	log.Info("Register> Hatchery %s registered with id:%d", h.Hatchery().Name, h.Hatchery().ID)

	if !uptodate {
		log.Warning("-=-=-=-=- Please update your hatchery binary -=-=-=-=-")
	}
	return nil
}

// GenerateName generate a hatchery's name
func GenerateName(add, name string) string {
	if name == "" {
		var errHostname error
		name, errHostname = os.Hostname()
		if errHostname != nil {
			log.Warning("Cannot retrieve hostname: %s", errHostname)
			name = "cds-hatchery"
		}
		name += "-" + namesgenerator.GetRandomName(0)
	}

	if add != "" {
		name += "-" + add
	}

	return name
}

func hearbeat(m Interface, token string, maxFailures int) {
	var failures int
	for {
		time.Sleep(5 * time.Second)
		if m.Hatchery().ID == 0 {
			log.Info("hearbeat> %s Disconnected from CDS engine, trying to register...", m.Hatchery().Name)
			if err := Register(m); err != nil {
				log.Info("hearbeat> %s Cannot register: %s", m.Hatchery().Name, err)
				checkFailures(maxFailures, failures)
				continue
			}
			if m.Hatchery().ID == 0 {
				log.Error("hearbeat> Cannot register hatchery. ID %d", m.Hatchery().ID)
				checkFailures(maxFailures, failures)
				continue
			}
			log.Info("hearbeat> %s Registered back: ID %d with model ID %d", m.Hatchery().Name, m.Hatchery().ID, m.Hatchery().Model.ID)
		}

		if err := m.Client().HatcheryRefresh(m.Hatchery().ID); err != nil {
			log.Info("heartbeat> %s cannot refresh beat: %s", m.Hatchery().Name, err)
			m.Hatchery().ID = 0
			checkFailures(maxFailures, failures)
			continue
		}
		failures = 0
	}
}

func checkFailures(maxFailures, nb int) {
	if nb > maxFailures {
		log.Error("Too many failures on try register. This hatchery is killed")
		os.Exit(10)
	}
}

func workerRegister(h Interface, models []sdk.Model) error {
	if len(models) == 0 {
		return fmt.Errorf("workerRegister> No model returned by GetWorkerModels")
	}
	log.Debug("workerRegister> models received: %d", len(models))

	var nRegistered int
	for _, m := range models {
		if m.Type != h.ModelType() {
			continue
		}

		// limit to 5 registration per ticker
		if nRegistered > 5 {
			break
		}
		// if current hatchery is in same group than worker model -> do not avoid spawn, even if worker model is in error
		if m.NbSpawnErr > 5 && h.Hatchery().GroupID != m.ID {
			log.Warning("workerRegister> Too many errors on spawn with model %s, please check this worker model", m.Name)
			continue
		}

		if h.NeedRegistration(&m) {
			log.Info("workerRegister> spawn a worker for register worker model %s (%d)", m.Name, m.ID)
			if _, errSpawn := h.SpawnWorker(&m, 0, nil, true, "spawn for register"); errSpawn != nil {
				log.Warning("workerRegister> cannot spawn worker for register: %s", m.Name, errSpawn)
				if err := h.Client().WorkerModelSpawnError(m.ID, fmt.Sprintf("workerRegister> cannot spawn worker for register: %s", errSpawn)); err != nil {
					log.Error("workerRegister> error on call client.WorkerModelSpawnError on worker model %s for register: %s", m.Name, errSpawn)
				}
				continue
			}
			nRegistered++
		} else {
			log.Debug("workerRegister> no need to register worker model %s (%d)", m.Name, m.ID)
			continue
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package game

import (
	"errors"
	"github.com/svera/acquire/game/board"
	"github.com/svera/acquire/game/corporation"
	"github.com/svera/acquire/game/player"
	"github.com/svera/acquire/game/tileset"
)

const totalCorporations = 7

type Game struct {
	board         *board.Board
	status        []string
	players       []*player.Player
	corporations  [7]*corporation.Corporation
	tileset       *tileset.Tileset
	currentPlayer uint
}

func New(
	board *board.Board, players []*player.Player, corporations [7]*corporation.Corporation, tileset *tileset.Tileset) (*Game, error) {
	if len(players) < 3 || len(players) > 6 {
		return nil, errors.New("Number of players must be between 3 and 6")
	}

	game := Game{
		board:         board,
		players:       players,
		corporations:  corporations,
		tileset:       tileset,
		currentPlayer: 0,
	}
	for _, player := range game.players {
		game.giveInitialTileset(player)
	}
	for i, corporation := range game.corporations {
		corporation.SetId(uint(i))
	}
	return &game, nil
}

func (g *Game) giveInitialTileset(player *player.Player) {
	for i := 0; i < 6; i++ {
		tile, _ := g.tileset.Draw()
		player.GetTile(tile)
	}
}

// Check if game end conditions are reached
func (g *Game) AreEndConditionsReached() bool {
	active := g.getActiveCorporations()
	if len(active) == 0 {
		return false
	}
	for _, corporation := range active {
		if corporation.Size() >= 41 {
			return true
		}
		if !corporation.IsSafe() {
			return false
		}
	}
	return true
}

// Return all corporations on the board
func (g *Game) getActiveCorporations() []*corporation.Corporation {
	active := []*corporation.Corporation{}
	for _, corporation := range g.corporations {
		if corporation.IsActive() {
			active = append(active, corporation)
		}
	}
	return active
}

// Placeholder function, pending implementation
func (g *Game) GetMainStockHolders(corporation *corporation.Corporation) [2][]*player.Player {
	mainStockHolders := [2][]*player.Player{{}, {}}
	stockHolders := g.getStockHolders(corporation)

	if len(stockHolders) == 1 {
		return [2][]*player.Player{{stockHolders[0]}, {stockHolders[0]}}
	}
	sharesDesc := func(p1, p2 *player.Player) bool {
		return p1.Shares(corporation) > p2.Shares(corporation)
	}
	player.By(sharesDesc).Sort(stockHolders)
	return mainStockHolders
}

// Get players who have stock of the passed corporation
func (g *Game) getStockHolders(corporation *corporation.Corporation) []*player.Player {
	var stockHolders []*player.Player

	for _, player := range g.players {
		if player.Shares(corporation) > 0 {
			stockHolders = append(stockHolders, player)
		}
	}
	return stockHolders
}

// Returns true if a tile is permanently unplayable, that is,
// that putting it on the board would merge two or more safe corporations
func (g *Game) isTileUnplayable(tile tileset.Position) bool {
	adjacents := g.board.AdjacentCells(tile)
	for _, adjacent := range adjacents {
		safeNeighbours := 0
		boardCell := g.board.Cell(adjacent)
		if boardCell != board.CellEmpty && boardCell != board.CellOrphanTile {
			if g.corporations[boardCell].IsSafe() {
				safeNeighbours++
			}
		}
		if safeNeighbours == 2 {
			return true
		}
	}
	return false
}

// Returns true if a tile is temporarily unplayable, that is,
// that putting it on the board would create an 8th corporation
func (g *Game) isTileTemporaryUnplayable(tile tileset.Position) bool {
	if len(g.getActiveCorporations()) < totalCorporations {
		return false
	}
	adjacents := g.board.AdjacentCells(tile)
	for _, adjacent := range adjacents {
		boardCell := g.board.Cell(adjacent)
		if boardCell == board.CellOrphanTile {
			return true
		}
	}
	return false
}
<commit_msg>Trying to make funcion more understandable<commit_after>package game

import (
	"errors"
	"github.com/svera/acquire/game/board"
	"github.com/svera/acquire/game/corporation"
	"github.com/svera/acquire/game/player"
	"github.com/svera/acquire/game/tileset"
)

const totalCorporations = 7

type Game struct {
	board         *board.Board
	status        []string
	players       []*player.Player
	corporations  [7]*corporation.Corporation
	tileset       *tileset.Tileset
	currentPlayer uint
}

func New(
	board *board.Board, players []*player.Player, corporations [7]*corporation.Corporation, tileset *tileset.Tileset) (*Game, error) {
	if len(players) < 3 || len(players) > 6 {
		return nil, errors.New("Number of players must be between 3 and 6")
	}

	game := Game{
		board:         board,
		players:       players,
		corporations:  corporations,
		tileset:       tileset,
		currentPlayer: 0,
	}
	for _, player := range game.players {
		game.giveInitialTileset(player)
	}
	for i, corporation := range game.corporations {
		corporation.SetId(uint(i))
	}
	return &game, nil
}

func (g *Game) giveInitialTileset(player *player.Player) {
	for i := 0; i < 6; i++ {
		tile, _ := g.tileset.Draw()
		player.GetTile(tile)
	}
}

// Check if game end conditions are reached
func (g *Game) AreEndConditionsReached() bool {
	active := g.getActiveCorporations()
	if len(active) == 0 {
		return false
	}
	for _, corporation := range active {
		if corporation.Size() >= 41 {
			return true
		}
		if !corporation.IsSafe() {
			return false
		}
	}
	return true
}

// Return all corporations on the board
func (g *Game) getActiveCorporations() []*corporation.Corporation {
	active := []*corporation.Corporation{}
	for _, corporation := range g.corporations {
		if corporation.IsActive() {
			active = append(active, corporation)
		}
	}
	return active
}

// Placeholder function, pending implementation
func (g *Game) GetMainStockHolders(corporation *corporation.Corporation) map[string][]*player.Player {
	mainStockHolders := map[string][]*player.Player{"primary": {}, "secondary": {}}
	stockHolders := g.getStockHolders(corporation)

	if len(stockHolders) == 1 {
		return map[string][]*player.Player{"primary": {stockHolders[0]}, "secondary": {stockHolders[0]}}
	}
	sharesDesc := func(p1, p2 *player.Player) bool {
		return p1.Shares(corporation) > p2.Shares(corporation)
	}
	player.By(sharesDesc).Sort(stockHolders)
	if stockHolders[0].Shares(corporation) == stockHolders[1].Shares(corporation) {
		mainStockHolders["primary"] = append(mainStockHolders["primary"], stockHolders[0])
		mainStockHolders["primary"] = append(mainStockHolders["primary"], stockHolders[1])
		i := 2
		for stockHolders[0] == stockHolders[i] && i < len(stockHolders) {
			mainStockHolders["primary"] = append(mainStockHolders["primary"], stockHolders[i])
			i++
		}
		return mainStockHolders
	}
	return mainStockHolders
}

// Get players who have stock of the passed corporation
func (g *Game) getStockHolders(corporation *corporation.Corporation) []*player.Player {
	var stockHolders []*player.Player

	for _, player := range g.players {
		if player.Shares(corporation) > 0 {
			stockHolders = append(stockHolders, player)
		}
	}
	return stockHolders
}

// Returns true if a tile is permanently unplayable, that is,
// that putting it on the board would merge two or more safe corporations
func (g *Game) isTileUnplayable(tile tileset.Position) bool {
	adjacents := g.board.AdjacentCells(tile)
	for _, adjacent := range adjacents {
		safeNeighbours := 0
		boardCell := g.board.Cell(adjacent)
		if boardCell != board.CellEmpty && boardCell != board.CellOrphanTile {
			if g.corporations[boardCell].IsSafe() {
				safeNeighbours++
			}
		}
		if safeNeighbours == 2 {
			return true
		}
	}
	return false
}

// Returns true if a tile is temporarily unplayable, that is,
// that putting it on the board would create an 8th corporation
func (g *Game) isTileTemporaryUnplayable(tile tileset.Position) bool {
	if len(g.getActiveCorporations()) < totalCorporations {
		return false
	}
	adjacents := g.board.AdjacentCells(tile)
	for _, adjacent := range adjacents {
		boardCell := g.board.Cell(adjacent)
		if boardCell == board.CellOrphanTile {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package gcsresource

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"

	"github.com/cheggaaa/pb"
	"golang.org/x/oauth2"
	oauthgoogle "golang.org/x/oauth2/google"
	"google.golang.org/api/storage/v1"
)

type GCSClient interface {
	BucketObjects(bucketName string, prefix string) ([]string, error)
	ObjectGenerations(bucketName string, objectPath string) ([]int64, error)
	DownloadFile(bucketName string, objectPath string, generation int64, localPath string) error
	UploadFile(bucketName string, objectPath string, localPath string) (int64, error)
	URL(bucketName string, objectPath string, generation int64) (string, error)
	DeleteObject(bucketName string, objectPath string, generation int64) error
}

type gcsclient struct {
	storageService  *storage.Service
	progressOutput io.Writer
}

func NewGCSClient(
	progressOutput io.Writer,
	jsonKey string,
) (GCSClient, error) {
	var err error
	var storageClient *http.Client
	var userAgent = "gcs-resource/0.0.1"

	if jsonKey != "" {
		storageJwtConf, err := oauthgoogle.JWTConfigFromJSON([]byte(jsonKey), storage.DevstorageFullControlScope)
		if err != nil {
			return &gcsclient{}, err
		}
		storageClient = storageJwtConf.Client(oauth2.NoContext)
	} else {
		storageClient, err = oauthgoogle.DefaultClient(oauth2.NoContext, storage.DevstorageFullControlScope)
		if err != nil {
			return &gcsclient{}, err
		}
	}

	storageService, err := storage.New(storageClient)
	if err != nil {
		return &gcsclient{}, err
	}
	storageService.UserAgent = userAgent

	return &gcsclient{
		storageService: storageService,
		progressOutput: progressOutput,
	}, nil
}

func (gcsclient *gcsclient) BucketObjects(bucketName string, prefix string) ([]string, error) {
	bucketObjects, err := gcsclient.getBucketObjects(bucketName, prefix)
	if err != nil {
		return []string{}, err
	}

	return bucketObjects, nil
}

func (gcsclient *gcsclient) ObjectGenerations(bucketName string, objectPath string) ([]int64, error) {
	isBucketVersioned, err := gcsclient.getBucketVersioning(bucketName)
	if err != nil {
		return []int64{}, err
	}

	if !isBucketVersioned {
		return []int64{}, errors.New("bucket is not versioned")
	}

	objectGenerations, err := gcsclient.getObjectGenerations(bucketName, objectPath)
	if err != nil {
		return []int64{}, err
	}

	return objectGenerations, nil
}

func (gcsclient *gcsclient) DownloadFile(bucketName string, objectPath string, generation int64, localPath string) error {
	isBucketVersioned, err := gcsclient.getBucketVersioning(bucketName)
	if err != nil {
		return err
	}

	if !isBucketVersioned && generation != 0 {
		return errors.New("bucket is not versioned")
	}

	getCall := gcsclient.storageService.Objects.Get(bucketName, objectPath)
	if generation != 0 {
		getCall = getCall.Generation(generation)
	}

	object, err := getCall.Do()
	if err != nil {
		return err
	}

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

	progress := gcsclient.newProgressBar(int64(object.Size))
	progress.Start()
	defer progress.Finish()

	response, err := getCall.Download()
	if err != nil {
		return err
	}
	defer response.Body.Close()

	reader := progress.NewProxyReader(response.Body)
	io.Copy(localFile, reader)

	return nil
}

func (gcsclient *gcsclient) UploadFile(bucketName string, objectPath string, localPath string) (int64, error) {
	isBucketVersioned, err := gcsclient.getBucketVersioning(bucketName)
	if err != nil {
		return 0, err
	}

	stat, err := os.Stat(localPath)
	if err != nil {
		return 0, err
	}

	localFile, err := os.Open(localPath)
	if err != nil {
		return 0, err
	}
	defer localFile.Close()

	progress := gcsclient.newProgressBar(stat.Size())
	progress.Start()
	defer progress.Finish()

	object := &storage.Object{
		Name: objectPath,
	}

	uploadedObject, err := gcsclient.storageService.Objects.Insert(bucketName, object).Media(progress.NewProxyReader(localFile)).Do()
	if err != nil {
		return 0, err
	}

	if isBucketVersioned {
		return uploadedObject.Generation, nil
	}

	return 0, nil
}

func (gcsclient *gcsclient) URL(bucketName string, objectPath string, generation int64) (string, error) {
	getCall := gcsclient.storageService.Objects.Get(bucketName, objectPath)
	if generation != 0 {
		getCall = getCall.Generation(generation)
	}

	_, err := getCall.Do()
	if err != nil {
		return "", err
	}

	var url string
	if generation != 0 {
		url = fmt.Sprintf("gs://%s/%s#%d", bucketName, objectPath, generation)
	} else {
		url = fmt.Sprintf("gs://%s/%s", bucketName, objectPath)
	}

	return url, nil
}

func (gcsclient *gcsclient) DeleteObject(bucketName string, objectPath string, generation int64) error {
	deleteCall := gcsclient.storageService.Objects.Delete(bucketName, objectPath)
	if generation != 0 {
		deleteCall = deleteCall.Generation(generation)
	}

	err := deleteCall.Do()
	if err != nil {
		return err
	}

	return nil
}

func (gcsclient *gcsclient) getBucketObjects(bucketName string, prefix string) ([]string, error) {
	var bucketObjects []string

	pageToken := ""
	for {
		listCall := gcsclient.storageService.Objects.List(bucketName)
		listCall = listCall.PageToken(pageToken)
		listCall = listCall.Prefix(prefix)
		listCall = listCall.Versions(false)

		objects, err := listCall.Do()
		if err != nil {
			return bucketObjects, err
		}

		for _, object := range objects.Items {
			bucketObjects = append(bucketObjects, object.Name)
		}

		if objects.NextPageToken != "" {
			pageToken = objects.NextPageToken
		} else {
			break
		}
	}

	return bucketObjects, nil
}

func (gcsclient *gcsclient) getBucketVersioning(bucketName string) (bool, error) {
	bucket, err := gcsclient.storageService.Buckets.Get(bucketName).Do()
	if err != nil {
		return false, err
	}

	return bucket.Versioning.Enabled, nil
}

func (gcsclient *gcsclient) getObjectGenerations(bucketName string, objectPath string) ([]int64, error) {
	var objectGenerations []int64

	pageToken := ""
	for {
		listCall := gcsclient.storageService.Objects.List(bucketName)
		listCall = listCall.PageToken(pageToken)
		listCall = listCall.Prefix(objectPath)
		listCall = listCall.Versions(true)

		objects, err := listCall.Do()
		if err != nil {
			return objectGenerations, err
		}

		for _, object := range objects.Items {
			if object.Name == objectPath {
				objectGenerations = append(objectGenerations, object.Generation)
			}
		}

		if objects.NextPageToken != "" {
			pageToken = objects.NextPageToken
		} else {
			break
		}
	}

	return objectGenerations, nil
}

func (gcsclient *gcsclient) newProgressBar(total int64) *pb.ProgressBar {
	progress := pb.New64(total)

	progress.Output = gcsclient.progressOutput
	progress.ShowSpeed = true
	progress.Units = pb.U_BYTES
	progress.NotPrint = true

	return progress.SetWidth(80)
}
<commit_msg>Fix checking if the bucket is versioned<commit_after>package gcsresource

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"

	"github.com/cheggaaa/pb"
	"golang.org/x/oauth2"
	oauthgoogle "golang.org/x/oauth2/google"
	"google.golang.org/api/storage/v1"
)

type GCSClient interface {
	BucketObjects(bucketName string, prefix string) ([]string, error)
	ObjectGenerations(bucketName string, objectPath string) ([]int64, error)
	DownloadFile(bucketName string, objectPath string, generation int64, localPath string) error
	UploadFile(bucketName string, objectPath string, localPath string) (int64, error)
	URL(bucketName string, objectPath string, generation int64) (string, error)
	DeleteObject(bucketName string, objectPath string, generation int64) error
}

type gcsclient struct {
	storageService  *storage.Service
	progressOutput io.Writer
}

func NewGCSClient(
	progressOutput io.Writer,
	jsonKey string,
) (GCSClient, error) {
	var err error
	var storageClient *http.Client
	var userAgent = "gcs-resource/0.0.1"

	if jsonKey != "" {
		storageJwtConf, err := oauthgoogle.JWTConfigFromJSON([]byte(jsonKey), storage.DevstorageFullControlScope)
		if err != nil {
			return &gcsclient{}, err
		}
		storageClient = storageJwtConf.Client(oauth2.NoContext)
	} else {
		storageClient, err = oauthgoogle.DefaultClient(oauth2.NoContext, storage.DevstorageFullControlScope)
		if err != nil {
			return &gcsclient{}, err
		}
	}

	storageService, err := storage.New(storageClient)
	if err != nil {
		return &gcsclient{}, err
	}
	storageService.UserAgent = userAgent

	return &gcsclient{
		storageService: storageService,
		progressOutput: progressOutput,
	}, nil
}

func (gcsclient *gcsclient) BucketObjects(bucketName string, prefix string) ([]string, error) {
	bucketObjects, err := gcsclient.getBucketObjects(bucketName, prefix)
	if err != nil {
		return []string{}, err
	}

	return bucketObjects, nil
}

func (gcsclient *gcsclient) ObjectGenerations(bucketName string, objectPath string) ([]int64, error) {
	isBucketVersioned, err := gcsclient.getBucketVersioning(bucketName)
	if err != nil {
		return []int64{}, err
	}

	if !isBucketVersioned {
		return []int64{}, errors.New("bucket is not versioned")
	}

	objectGenerations, err := gcsclient.getObjectGenerations(bucketName, objectPath)
	if err != nil {
		return []int64{}, err
	}

	return objectGenerations, nil
}

func (gcsclient *gcsclient) DownloadFile(bucketName string, objectPath string, generation int64, localPath string) error {
	isBucketVersioned, err := gcsclient.getBucketVersioning(bucketName)
	if err != nil {
		return err
	}

	if !isBucketVersioned && generation != 0 {
		return errors.New("bucket is not versioned")
	}

	getCall := gcsclient.storageService.Objects.Get(bucketName, objectPath)
	if generation != 0 {
		getCall = getCall.Generation(generation)
	}

	object, err := getCall.Do()
	if err != nil {
		return err
	}

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

	progress := gcsclient.newProgressBar(int64(object.Size))
	progress.Start()
	defer progress.Finish()

	response, err := getCall.Download()
	if err != nil {
		return err
	}
	defer response.Body.Close()

	reader := progress.NewProxyReader(response.Body)
	io.Copy(localFile, reader)

	return nil
}

func (gcsclient *gcsclient) UploadFile(bucketName string, objectPath string, localPath string) (int64, error) {
	isBucketVersioned, err := gcsclient.getBucketVersioning(bucketName)
	if err != nil {
		return 0, err
	}

	stat, err := os.Stat(localPath)
	if err != nil {
		return 0, err
	}

	localFile, err := os.Open(localPath)
	if err != nil {
		return 0, err
	}
	defer localFile.Close()

	progress := gcsclient.newProgressBar(stat.Size())
	progress.Start()
	defer progress.Finish()

	object := &storage.Object{
		Name: objectPath,
	}

	uploadedObject, err := gcsclient.storageService.Objects.Insert(bucketName, object).Media(progress.NewProxyReader(localFile)).Do()
	if err != nil {
		return 0, err
	}

	if isBucketVersioned {
		return uploadedObject.Generation, nil
	}

	return 0, nil
}

func (gcsclient *gcsclient) URL(bucketName string, objectPath string, generation int64) (string, error) {
	getCall := gcsclient.storageService.Objects.Get(bucketName, objectPath)
	if generation != 0 {
		getCall = getCall.Generation(generation)
	}

	_, err := getCall.Do()
	if err != nil {
		return "", err
	}

	var url string
	if generation != 0 {
		url = fmt.Sprintf("gs://%s/%s#%d", bucketName, objectPath, generation)
	} else {
		url = fmt.Sprintf("gs://%s/%s", bucketName, objectPath)
	}

	return url, nil
}

func (gcsclient *gcsclient) DeleteObject(bucketName string, objectPath string, generation int64) error {
	deleteCall := gcsclient.storageService.Objects.Delete(bucketName, objectPath)
	if generation != 0 {
		deleteCall = deleteCall.Generation(generation)
	}

	err := deleteCall.Do()
	if err != nil {
		return err
	}

	return nil
}

func (gcsclient *gcsclient) getBucketObjects(bucketName string, prefix string) ([]string, error) {
	var bucketObjects []string

	pageToken := ""
	for {
		listCall := gcsclient.storageService.Objects.List(bucketName)
		listCall = listCall.PageToken(pageToken)
		listCall = listCall.Prefix(prefix)
		listCall = listCall.Versions(false)

		objects, err := listCall.Do()
		if err != nil {
			return bucketObjects, err
		}

		for _, object := range objects.Items {
			bucketObjects = append(bucketObjects, object.Name)
		}

		if objects.NextPageToken != "" {
			pageToken = objects.NextPageToken
		} else {
			break
		}
	}

	return bucketObjects, nil
}

func (gcsclient *gcsclient) getBucketVersioning(bucketName string) (bool, error) {
	bucket, err := gcsclient.storageService.Buckets.Get(bucketName).Do()
	if err != nil {
		return false, err
	}

	if bucket.Versioning != nil {
		return bucket.Versioning.Enabled, nil
	}

	return false, nil
}

func (gcsclient *gcsclient) getObjectGenerations(bucketName string, objectPath string) ([]int64, error) {
	var objectGenerations []int64

	pageToken := ""
	for {
		listCall := gcsclient.storageService.Objects.List(bucketName)
		listCall = listCall.PageToken(pageToken)
		listCall = listCall.Prefix(objectPath)
		listCall = listCall.Versions(true)

		objects, err := listCall.Do()
		if err != nil {
			return objectGenerations, err
		}

		for _, object := range objects.Items {
			if object.Name == objectPath {
				objectGenerations = append(objectGenerations, object.Generation)
			}
		}

		if objects.NextPageToken != "" {
			pageToken = objects.NextPageToken
		} else {
			break
		}
	}

	return objectGenerations, nil
}

func (gcsclient *gcsclient) newProgressBar(total int64) *pb.ProgressBar {
	progress := pb.New64(total)

	progress.Output = gcsclient.progressOutput
	progress.ShowSpeed = true
	progress.Units = pb.U_BYTES
	progress.NotPrint = true

	return progress.SetWidth(80)
}
<|endoftext|>


<commit_before>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// Package api Golang bindings for Baruwa REST API
package api

import (
	"fmt"
	"net/url"

	"github.com/google/go-querystring/query"
)

// Alias hosts alias addresses
// https://www.baruwa.com/docs/api/#alias-addresses
type Alias struct {
	ID      int    `json:"id,omitempty"`
	Address string `json:"address"`
	Enabled bool   `json:"enabled"`
}

// GetAlias returns an alias address
// https://www.baruwa.com/docs/api/#retrieve-an-existing-alias-address
func (c *Client) GetAlias(id int) (alias *Alias, err error) {
	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	err = c.get(fmt.Sprintf("aliasaddresses/%d", id), alias)

	return
}

// CreateAlias creates an alias address
// https://www.baruwa.com/docs/api/#create-an-alias-address
func (c *Client) CreateAlias(userid int, alias *Alias) (err error) {
	var v url.Values

	if userid <= 0 {
		err = fmt.Errorf("The userid param should be > 0")
		return
	}

	if alias == nil {
		err = fmt.Errorf("The alias param cannot be nil")
		return
	}

	if v, err = query.Values(alias); err != nil {
		return
	}

	err = c.post(fmt.Sprintf("aliasaddresses/%d", userid), v, alias)

	return
}

// UpdateAlias updates an alias address
// https://www.baruwa.com/docs/api/#update-an-alias-address
func (c *Client) UpdateAlias(alias *Alias) (err error) {
	var v url.Values

	if alias == nil {
		err = fmt.Errorf("The alias param cannot be nil")
		return
	}

	if alias.ID <= 0 {
		err = fmt.Errorf("The alias.ID param should be > 0")
		return
	}

	if v, err = query.Values(alias); err != nil {
		return
	}

	err = c.put(fmt.Sprintf("aliasaddresses/%d", alias.ID), v, alias)

	return
}

// DeleteAlias deletes an alias address
// https://www.baruwa.com/docs/api/#delete-an-alias-address
func (c *Client) DeleteAlias(id int) (err error) {
	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	err = c.delete(fmt.Sprintf("aliasaddresses/%d", id))

	return
}
<commit_msg>FIX: Rename alias address struct<commit_after>// BaruwaAPI Golang bindings for Baruwa REST API
// Copyright (C) 2019 Andrew Colin Kissa <andrew@topdog.za.net>

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at http://mozilla.org/MPL/2.0/.

// Package api Golang bindings for Baruwa REST API
package api

import (
	"fmt"
	"net/url"

	"github.com/google/go-querystring/query"
)

// AliasAddress hosts alias addresses
// https://www.baruwa.com/docs/api/#alias-addresses
type AliasAddress struct {
	ID      int    `json:"id,omitempty"`
	Address string `json:"address"`
	Enabled bool   `json:"enabled"`
}

// GetAliasAddress returns an alias address
// https://www.baruwa.com/docs/api/#retrieve-an-existing-alias-address
func (c *Client) GetAliasAddress(id int) (alias *AliasAddress, err error) {
	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	err = c.get(fmt.Sprintf("aliasaddresses/%d", id), alias)

	return
}

// CreateAliasAddress creates an alias address
// https://www.baruwa.com/docs/api/#create-an-alias-address
func (c *Client) CreateAliasAddress(userid int, alias *AliasAddress) (err error) {
	var v url.Values

	if userid <= 0 {
		err = fmt.Errorf("The userid param should be > 0")
		return
	}

	if alias == nil {
		err = fmt.Errorf("The alias param cannot be nil")
		return
	}

	if v, err = query.Values(alias); err != nil {
		return
	}

	err = c.post(fmt.Sprintf("aliasaddresses/%d", userid), v, alias)

	return
}

// UpdateAliasAddress updates an alias address
// https://www.baruwa.com/docs/api/#update-an-alias-address
func (c *Client) UpdateAliasAddress(alias *AliasAddress) (err error) {
	var v url.Values

	if alias == nil {
		err = fmt.Errorf("The alias param cannot be nil")
		return
	}

	if alias.ID <= 0 {
		err = fmt.Errorf("The alias.ID param should be > 0")
		return
	}

	if v, err = query.Values(alias); err != nil {
		return
	}

	err = c.put(fmt.Sprintf("aliasaddresses/%d", alias.ID), v, alias)

	return
}

// DeleteAliasAddress deletes an alias address
// https://www.baruwa.com/docs/api/#delete-an-alias-address
func (c *Client) DeleteAliasAddress(id int) (err error) {
	if id <= 0 {
		err = fmt.Errorf("The id param should be > 0")
		return
	}

	err = c.delete(fmt.Sprintf("aliasaddresses/%d", id))

	return
}
<|endoftext|>


<commit_before>package image // import "github.com/docker/docker/api/server/router/image"

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"strconv"
	"strings"

	"github.com/containerd/containerd/platforms"
	"github.com/docker/docker/api/server/httputils"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/versions"
	"github.com/docker/docker/errdefs"
	"github.com/docker/docker/pkg/ioutils"
	"github.com/docker/docker/pkg/streamformatter"
	"github.com/docker/docker/registry"
	specs "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

// Creates an image from Pull or from Import
func (s *imageRouter) postImagesCreate(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {

	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	var (
		image    = r.Form.Get("fromImage")
		repo     = r.Form.Get("repo")
		tag      = r.Form.Get("tag")
		message  = r.Form.Get("message")
		err      error
		output   = ioutils.NewWriteFlusher(w)
		platform *specs.Platform
	)
	defer output.Close()

	w.Header().Set("Content-Type", "application/json")

	version := httputils.VersionFromContext(ctx)
	if versions.GreaterThanOrEqualTo(version, "1.32") {
		apiPlatform := r.FormValue("platform")
		if apiPlatform != "" {
			sp, err := platforms.Parse(apiPlatform)
			if err != nil {
				return err
			}
			platform = &sp
		}
	}

	if image != "" { // pull
		metaHeaders := map[string][]string{}
		for k, v := range r.Header {
			if strings.HasPrefix(k, "X-Meta-") {
				metaHeaders[k] = v
			}
		}

		authEncoded := r.Header.Get("X-Registry-Auth")
		authConfig := &types.AuthConfig{}
		if authEncoded != "" {
			authJSON := base64.NewDecoder(base64.URLEncoding, strings.NewReader(authEncoded))
			if err := json.NewDecoder(authJSON).Decode(authConfig); err != nil {
				// for a pull it is not an error if no auth was given
				// to increase compatibility with the existing api it is defaulting to be empty
				authConfig = &types.AuthConfig{}
			}
		}
		err = s.backend.PullImage(ctx, image, tag, platform, metaHeaders, authConfig, output)
	} else { // import
		src := r.Form.Get("fromSrc")
		// 'err' MUST NOT be defined within this block, we need any error
		// generated from the download to be available to the output
		// stream processing below
		os := ""
		if platform != nil {
			os = platform.OS
		}
		err = s.backend.ImportImage(src, repo, os, tag, message, r.Body, output, r.Form["changes"])
	}
	if err != nil {
		if !output.Flushed() {
			return err
		}
		_, _ = output.Write(streamformatter.FormatError(err))
	}

	return nil
}

func (s *imageRouter) postImagesPush(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	metaHeaders := map[string][]string{}
	for k, v := range r.Header {
		if strings.HasPrefix(k, "X-Meta-") {
			metaHeaders[k] = v
		}
	}
	if err := httputils.ParseForm(r); err != nil {
		return err
	}
	authConfig := &types.AuthConfig{}

	authEncoded := r.Header.Get("X-Registry-Auth")
	if authEncoded != "" {
		// the new format is to handle the authConfig as a header
		authJSON := base64.NewDecoder(base64.URLEncoding, strings.NewReader(authEncoded))
		if err := json.NewDecoder(authJSON).Decode(authConfig); err != nil {
			// to increase compatibility to existing api it is defaulting to be empty
			authConfig = &types.AuthConfig{}
		}
	} else {
		// the old format is supported for compatibility if there was no authConfig header
		if err := json.NewDecoder(r.Body).Decode(authConfig); err != nil {
			return errors.Wrap(errdefs.InvalidParameter(err), "Bad parameters and missing X-Registry-Auth")
		}
	}

	image := vars["name"]
	tag := r.Form.Get("tag")

	output := ioutils.NewWriteFlusher(w)
	defer output.Close()

	w.Header().Set("Content-Type", "application/json")

	if err := s.backend.PushImage(ctx, image, tag, metaHeaders, authConfig, output); err != nil {
		if !output.Flushed() {
			return err
		}
		_, _ = output.Write(streamformatter.FormatError(err))
	}
	return nil
}

func (s *imageRouter) getImagesGet(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	w.Header().Set("Content-Type", "application/x-tar")

	output := ioutils.NewWriteFlusher(w)
	defer output.Close()
	var names []string
	if name, ok := vars["name"]; ok {
		names = []string{name}
	} else {
		names = r.Form["names"]
	}

	if err := s.backend.ExportImage(names, output); err != nil {
		if !output.Flushed() {
			return err
		}
		_, _ = output.Write(streamformatter.FormatError(err))
	}
	return nil
}

func (s *imageRouter) postImagesLoad(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}
	quiet := httputils.BoolValueOrDefault(r, "quiet", true)

	w.Header().Set("Content-Type", "application/json")

	output := ioutils.NewWriteFlusher(w)
	defer output.Close()
	if err := s.backend.LoadImage(r.Body, output, quiet); err != nil {
		_, _ = output.Write(streamformatter.FormatError(err))
	}
	return nil
}

type missingImageError struct{}

func (missingImageError) Error() string {
	return "image name cannot be blank"
}

func (missingImageError) InvalidParameter() {}

func (s *imageRouter) deleteImages(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	name := vars["name"]

	if strings.TrimSpace(name) == "" {
		return missingImageError{}
	}

	force := httputils.BoolValue(r, "force")
	prune := !httputils.BoolValue(r, "noprune")

	list, err := s.backend.ImageDelete(name, force, prune)
	if err != nil {
		return err
	}

	return httputils.WriteJSON(w, http.StatusOK, list)
}

func (s *imageRouter) getImagesByName(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	imageInspect, err := s.backend.LookupImage(vars["name"])
	if err != nil {
		return err
	}

	return httputils.WriteJSON(w, http.StatusOK, imageInspect)
}

func (s *imageRouter) getImagesJSON(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	imageFilters, err := filters.FromJSON(r.Form.Get("filters"))
	if err != nil {
		return err
	}

	version := httputils.VersionFromContext(ctx)
	if versions.LessThan(version, "1.41") {
		// NOTE: filter is a shell glob string applied to repository names.
		filterParam := r.Form.Get("filter")
		if filterParam != "" {
			imageFilters.Add("reference", filterParam)
		}
	}

	var sharedSize bool
	if versions.GreaterThanOrEqualTo(version, "1.42") {
		// NOTE: Support for the "shared-size" parameter was added in API 1.42.
		sharedSize = httputils.BoolValue(r, "shared-size")
	}

	images, err := s.backend.Images(ctx, types.ImageListOptions{
		All:        httputils.BoolValue(r, "all"),
		Filters:    imageFilters,
		SharedSize: sharedSize,
	})
	if err != nil {
		return err
	}

	return httputils.WriteJSON(w, http.StatusOK, images)
}

func (s *imageRouter) getImagesHistory(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	name := vars["name"]
	history, err := s.backend.ImageHistory(name)
	if err != nil {
		return err
	}

	return httputils.WriteJSON(w, http.StatusOK, history)
}

func (s *imageRouter) postImagesTag(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}
	if _, err := s.backend.TagImage(vars["name"], r.Form.Get("repo"), r.Form.Get("tag")); err != nil {
		return err
	}
	w.WriteHeader(http.StatusCreated)
	return nil
}

func (s *imageRouter) getImagesSearch(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}
	var (
		config      *types.AuthConfig
		authEncoded = r.Header.Get("X-Registry-Auth")
		headers     = map[string][]string{}
	)

	if authEncoded != "" {
		authJSON := base64.NewDecoder(base64.URLEncoding, strings.NewReader(authEncoded))
		if err := json.NewDecoder(authJSON).Decode(&config); err != nil {
			// for a search it is not an error if no auth was given
			// to increase compatibility with the existing api it is defaulting to be empty
			config = &types.AuthConfig{}
		}
	}
	for k, v := range r.Header {
		if strings.HasPrefix(k, "X-Meta-") {
			headers[k] = v
		}
	}
	limit := registry.DefaultSearchLimit
	if r.Form.Get("limit") != "" {
		limitValue, err := strconv.Atoi(r.Form.Get("limit"))
		if err != nil {
			return err
		}
		limit = limitValue
	}
	query, err := s.backend.SearchRegistryForImages(ctx, r.Form.Get("filters"), r.Form.Get("term"), limit, config, headers)
	if err != nil {
		return err
	}
	return httputils.WriteJSON(w, http.StatusOK, query.Results)
}

func (s *imageRouter) postImagesPrune(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	pruneFilters, err := filters.FromJSON(r.Form.Get("filters"))
	if err != nil {
		return err
	}

	pruneReport, err := s.backend.ImagesPrune(ctx, pruneFilters)
	if err != nil {
		return err
	}
	return httputils.WriteJSON(w, http.StatusOK, pruneReport)
}
<commit_msg>api: postImagesCreate(): rename ambiguous err variable<commit_after>package image // import "github.com/docker/docker/api/server/router/image"

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"strconv"
	"strings"

	"github.com/containerd/containerd/platforms"
	"github.com/docker/docker/api/server/httputils"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	"github.com/docker/docker/api/types/versions"
	"github.com/docker/docker/errdefs"
	"github.com/docker/docker/pkg/ioutils"
	"github.com/docker/docker/pkg/streamformatter"
	"github.com/docker/docker/registry"
	specs "github.com/opencontainers/image-spec/specs-go/v1"
	"github.com/pkg/errors"
)

// Creates an image from Pull or from Import
func (s *imageRouter) postImagesCreate(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {

	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	var (
		image       = r.Form.Get("fromImage")
		repo        = r.Form.Get("repo")
		tag         = r.Form.Get("tag")
		message     = r.Form.Get("message")
		progressErr error
		output      = ioutils.NewWriteFlusher(w)
		platform    *specs.Platform
	)
	defer output.Close()

	w.Header().Set("Content-Type", "application/json")

	version := httputils.VersionFromContext(ctx)
	if versions.GreaterThanOrEqualTo(version, "1.32") {
		apiPlatform := r.FormValue("platform")
		if apiPlatform != "" {
			sp, err := platforms.Parse(apiPlatform)
			if err != nil {
				return err
			}
			platform = &sp
		}
	}

	if image != "" { // pull
		metaHeaders := map[string][]string{}
		for k, v := range r.Header {
			if strings.HasPrefix(k, "X-Meta-") {
				metaHeaders[k] = v
			}
		}

		authEncoded := r.Header.Get("X-Registry-Auth")
		authConfig := &types.AuthConfig{}
		if authEncoded != "" {
			authJSON := base64.NewDecoder(base64.URLEncoding, strings.NewReader(authEncoded))
			if err := json.NewDecoder(authJSON).Decode(authConfig); err != nil {
				// for a pull it is not an error if no auth was given
				// to increase compatibility with the existing api it is defaulting to be empty
				authConfig = &types.AuthConfig{}
			}
		}
		progressErr = s.backend.PullImage(ctx, image, tag, platform, metaHeaders, authConfig, output)
	} else { // import
		src := r.Form.Get("fromSrc")
		os := ""
		if platform != nil {
			os = platform.OS
		}
		progressErr = s.backend.ImportImage(src, repo, os, tag, message, r.Body, output, r.Form["changes"])
	}
	if progressErr != nil {
		if !output.Flushed() {
			return progressErr
		}
		_, _ = output.Write(streamformatter.FormatError(progressErr))
	}

	return nil
}

func (s *imageRouter) postImagesPush(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	metaHeaders := map[string][]string{}
	for k, v := range r.Header {
		if strings.HasPrefix(k, "X-Meta-") {
			metaHeaders[k] = v
		}
	}
	if err := httputils.ParseForm(r); err != nil {
		return err
	}
	authConfig := &types.AuthConfig{}

	authEncoded := r.Header.Get("X-Registry-Auth")
	if authEncoded != "" {
		// the new format is to handle the authConfig as a header
		authJSON := base64.NewDecoder(base64.URLEncoding, strings.NewReader(authEncoded))
		if err := json.NewDecoder(authJSON).Decode(authConfig); err != nil {
			// to increase compatibility to existing api it is defaulting to be empty
			authConfig = &types.AuthConfig{}
		}
	} else {
		// the old format is supported for compatibility if there was no authConfig header
		if err := json.NewDecoder(r.Body).Decode(authConfig); err != nil {
			return errors.Wrap(errdefs.InvalidParameter(err), "Bad parameters and missing X-Registry-Auth")
		}
	}

	image := vars["name"]
	tag := r.Form.Get("tag")

	output := ioutils.NewWriteFlusher(w)
	defer output.Close()

	w.Header().Set("Content-Type", "application/json")

	if err := s.backend.PushImage(ctx, image, tag, metaHeaders, authConfig, output); err != nil {
		if !output.Flushed() {
			return err
		}
		_, _ = output.Write(streamformatter.FormatError(err))
	}
	return nil
}

func (s *imageRouter) getImagesGet(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	w.Header().Set("Content-Type", "application/x-tar")

	output := ioutils.NewWriteFlusher(w)
	defer output.Close()
	var names []string
	if name, ok := vars["name"]; ok {
		names = []string{name}
	} else {
		names = r.Form["names"]
	}

	if err := s.backend.ExportImage(names, output); err != nil {
		if !output.Flushed() {
			return err
		}
		_, _ = output.Write(streamformatter.FormatError(err))
	}
	return nil
}

func (s *imageRouter) postImagesLoad(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}
	quiet := httputils.BoolValueOrDefault(r, "quiet", true)

	w.Header().Set("Content-Type", "application/json")

	output := ioutils.NewWriteFlusher(w)
	defer output.Close()
	if err := s.backend.LoadImage(r.Body, output, quiet); err != nil {
		_, _ = output.Write(streamformatter.FormatError(err))
	}
	return nil
}

type missingImageError struct{}

func (missingImageError) Error() string {
	return "image name cannot be blank"
}

func (missingImageError) InvalidParameter() {}

func (s *imageRouter) deleteImages(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	name := vars["name"]

	if strings.TrimSpace(name) == "" {
		return missingImageError{}
	}

	force := httputils.BoolValue(r, "force")
	prune := !httputils.BoolValue(r, "noprune")

	list, err := s.backend.ImageDelete(name, force, prune)
	if err != nil {
		return err
	}

	return httputils.WriteJSON(w, http.StatusOK, list)
}

func (s *imageRouter) getImagesByName(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	imageInspect, err := s.backend.LookupImage(vars["name"])
	if err != nil {
		return err
	}

	return httputils.WriteJSON(w, http.StatusOK, imageInspect)
}

func (s *imageRouter) getImagesJSON(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	imageFilters, err := filters.FromJSON(r.Form.Get("filters"))
	if err != nil {
		return err
	}

	version := httputils.VersionFromContext(ctx)
	if versions.LessThan(version, "1.41") {
		// NOTE: filter is a shell glob string applied to repository names.
		filterParam := r.Form.Get("filter")
		if filterParam != "" {
			imageFilters.Add("reference", filterParam)
		}
	}

	var sharedSize bool
	if versions.GreaterThanOrEqualTo(version, "1.42") {
		// NOTE: Support for the "shared-size" parameter was added in API 1.42.
		sharedSize = httputils.BoolValue(r, "shared-size")
	}

	images, err := s.backend.Images(ctx, types.ImageListOptions{
		All:        httputils.BoolValue(r, "all"),
		Filters:    imageFilters,
		SharedSize: sharedSize,
	})
	if err != nil {
		return err
	}

	return httputils.WriteJSON(w, http.StatusOK, images)
}

func (s *imageRouter) getImagesHistory(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	name := vars["name"]
	history, err := s.backend.ImageHistory(name)
	if err != nil {
		return err
	}

	return httputils.WriteJSON(w, http.StatusOK, history)
}

func (s *imageRouter) postImagesTag(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}
	if _, err := s.backend.TagImage(vars["name"], r.Form.Get("repo"), r.Form.Get("tag")); err != nil {
		return err
	}
	w.WriteHeader(http.StatusCreated)
	return nil
}

func (s *imageRouter) getImagesSearch(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}
	var (
		config      *types.AuthConfig
		authEncoded = r.Header.Get("X-Registry-Auth")
		headers     = map[string][]string{}
	)

	if authEncoded != "" {
		authJSON := base64.NewDecoder(base64.URLEncoding, strings.NewReader(authEncoded))
		if err := json.NewDecoder(authJSON).Decode(&config); err != nil {
			// for a search it is not an error if no auth was given
			// to increase compatibility with the existing api it is defaulting to be empty
			config = &types.AuthConfig{}
		}
	}
	for k, v := range r.Header {
		if strings.HasPrefix(k, "X-Meta-") {
			headers[k] = v
		}
	}
	limit := registry.DefaultSearchLimit
	if r.Form.Get("limit") != "" {
		limitValue, err := strconv.Atoi(r.Form.Get("limit"))
		if err != nil {
			return err
		}
		limit = limitValue
	}
	query, err := s.backend.SearchRegistryForImages(ctx, r.Form.Get("filters"), r.Form.Get("term"), limit, config, headers)
	if err != nil {
		return err
	}
	return httputils.WriteJSON(w, http.StatusOK, query.Results)
}

func (s *imageRouter) postImagesPrune(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
	if err := httputils.ParseForm(r); err != nil {
		return err
	}

	pruneFilters, err := filters.FromJSON(r.Form.Get("filters"))
	if err != nil {
		return err
	}

	pruneReport, err := s.backend.ImagesPrune(ctx, pruneFilters)
	if err != nil {
		return err
	}
	return httputils.WriteJSON(w, http.StatusOK, pruneReport)
}
<|endoftext|>


<commit_before>package agent

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"github.com/buildkite/agent/logger"
	"golang.org/x/oauth2/google"
	storage "google.golang.org/api/storage/v1"
)

type GSDownloaderConfig struct {
	// The name of the bucket
	Bucket string

	// The root directory of the download
	Destination string

	// The relative path that should be preserved in the download folder,
	// also it's location in the bucket
	Path string

	// How many times should it retry the download before giving up
	Retries int

	// If failed responses should be dumped to the log
	DebugHTTP bool
}

type GSDownloader struct {
	// The config for the downloader
	conf GSDownloaderConfig

	// The logger instance to use
	logger logger.Logger
}

func NewGSDownloader(l logger.Logger, c GSDownloaderConfig) *GSDownloader {
	return &GSDownloader{
		logger: l,
		conf:   c,
	}
}

func (d GSDownloader) Start() error {
	client, err := google.DefaultClient(context.Background(), storage.DevstorageReadOnlyScope)
	if err != nil {
		return errors.New(fmt.Sprintf("Error creating Google Cloud Storage client: %v", err))
	}

	url := "https://www.googleapis.com/storage/v1/b/" + d.BucketName() + "/o/" + escape(d.BucketFileLocation()) + "?alt=media"

	// We can now cheat and pass the URL onto our regular downloader
	return NewDownload(d.logger, client, DownloadConfig{
		URL:         url,
		Path:        d.conf.Path,
		Destination: d.conf.Destination,
		Retries:     d.conf.Retries,
		DebugHTTP:   d.conf.DebugHTTP,
	}).Start()
}

func (d GSDownloader) BucketFileLocation() string {
	if d.BucketPath() != "" {
		return strings.TrimSuffix(d.BucketPath(), "/") + "/" + strings.TrimPrefix(d.conf.Path, "/")
	} else {
		return d.conf.Path
	}
}

func (d GSDownloader) BucketPath() string {
	return strings.Join(d.destinationParts()[1:len(d.destinationParts())], "/")
}

func (d GSDownloader) BucketName() string {
	return d.destinationParts()[0]
}

func (d GSDownloader) destinationParts() []string {
	trimmed := strings.TrimPrefix(d.conf.Bucket, "gs://")

	return strings.Split(trimmed, "/")
}

func escape(s string) string {
	// See https://golang.org/src/net/url/url.go
	hexCount := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		if shouldEscape(c) {
			hexCount++
		}
	}

	if hexCount == 0 {
		return s
	}

	t := make([]byte, len(s)+2*hexCount)
	j := 0
	for i := 0; i < len(s); i++ {
		switch c := s[i]; {
		case shouldEscape(c):
			t[j] = '%'
			t[j+1] = "0123456789ABCDEF"[c>>4]
			t[j+2] = "0123456789ABCDEF"[c&15]
			j += 3
		default:
			t[j] = s[i]
			j++
		}
	}
	return string(t)
}

func shouldEscape(c byte) bool {
	// See https://cloud.google.com/storage/docs/json_api/#encoding
	if 'A' <= c && c <= 'Z' || 'a' <= c && c <= 'z' || '0' <= c && c <= '9' {
		return false
	}
	switch c {
	case '-', '.', '_', '~', '!', '$', '&', '\'', '(', ')', '*', '+', ',', ';', '=', ':', '@':
		return false
	}
	return true
}
<commit_msg>allow downloader to use GS_APPLICATION_CREDENTIALS<commit_after>package agent

import (
	"errors"
	"fmt"
	"strings"

	"github.com/buildkite/agent/logger"
	storage "google.golang.org/api/storage/v1"
)

type GSDownloaderConfig struct {
	// The name of the bucket
	Bucket string

	// The root directory of the download
	Destination string

	// The relative path that should be preserved in the download folder,
	// also it's location in the bucket
	Path string

	// How many times should it retry the download before giving up
	Retries int

	// If failed responses should be dumped to the log
	DebugHTTP bool
}

type GSDownloader struct {
	// The config for the downloader
	conf GSDownloaderConfig

	// The logger instance to use
	logger logger.Logger
}

func NewGSDownloader(l logger.Logger, c GSDownloaderConfig) *GSDownloader {
	return &GSDownloader{
		logger: l,
		conf:   c,
	}
}

func (d GSDownloader) Start() error {
	client, err := newGoogleClient(storage.DevstorageReadOnlyScope)
	if err != nil {
		return errors.New(fmt.Sprintf("Error creating Google Cloud Storage client: %v", err))
	}

	url := "https://www.googleapis.com/storage/v1/b/" + d.BucketName() + "/o/" + escape(d.BucketFileLocation()) + "?alt=media"

	// We can now cheat and pass the URL onto our regular downloader
	return NewDownload(d.logger, client, DownloadConfig{
		URL:         url,
		Path:        d.conf.Path,
		Destination: d.conf.Destination,
		Retries:     d.conf.Retries,
		DebugHTTP:   d.conf.DebugHTTP,
	}).Start()
}

func (d GSDownloader) BucketFileLocation() string {
	if d.BucketPath() != "" {
		return strings.TrimSuffix(d.BucketPath(), "/") + "/" + strings.TrimPrefix(d.conf.Path, "/")
	} else {
		return d.conf.Path
	}
}

func (d GSDownloader) BucketPath() string {
	return strings.Join(d.destinationParts()[1:len(d.destinationParts())], "/")
}

func (d GSDownloader) BucketName() string {
	return d.destinationParts()[0]
}

func (d GSDownloader) destinationParts() []string {
	trimmed := strings.TrimPrefix(d.conf.Bucket, "gs://")

	return strings.Split(trimmed, "/")
}

func escape(s string) string {
	// See https://golang.org/src/net/url/url.go
	hexCount := 0
	for i := 0; i < len(s); i++ {
		c := s[i]
		if shouldEscape(c) {
			hexCount++
		}
	}

	if hexCount == 0 {
		return s
	}

	t := make([]byte, len(s)+2*hexCount)
	j := 0
	for i := 0; i < len(s); i++ {
		switch c := s[i]; {
		case shouldEscape(c):
			t[j] = '%'
			t[j+1] = "0123456789ABCDEF"[c>>4]
			t[j+2] = "0123456789ABCDEF"[c&15]
			j += 3
		default:
			t[j] = s[i]
			j++
		}
	}
	return string(t)
}

func shouldEscape(c byte) bool {
	// See https://cloud.google.com/storage/docs/json_api/#encoding
	if 'A' <= c && c <= 'Z' || 'a' <= c && c <= 'z' || '0' <= c && c <= '9' {
		return false
	}
	switch c {
	case '-', '.', '_', '~', '!', '$', '&', '\'', '(', ')', '*', '+', ',', ';', '=', ':', '@':
		return false
	}
	return true
}
<|endoftext|>


<commit_before>package bolt

import (
	"sync"
	"testing"
)

const (
	BenchReadMode            = "read"
	BenchWriteMode           = "write"
	BenchSequentialTraversal = "sequential"
	BenchRandomTraversal     = "random"
)

type Benchmark struct {
	InputPath        string
	ReadWriteMode    string
	TraversalPattern string
	Parallelism      int
}

func NewBenchmark(inputPath, readWriteMode, traversalPattern string, parallelism int) *Benchmark {
	return &Benchmark{inputPath, readWriteMode, traversalPattern, parallelism}
}

func (bm *Benchmark) Run(b *testing.B) {

	// Open the database.
	db, err := Open(bm.InputPath, 0600)
	if err != nil {
		panic(err)
		return
	}
	defer db.Close()

	b.ResetTimer()

	// Keep running a fixed number of parallel reads until we run out of time.
	for i := 0; i < b.N; i++ {
		var wg sync.WaitGroup
		for j := 0; j < bm.Parallelism; j++ {
			wg.Add(1)
			go func() {
				if bm.TraversalPattern == BenchRandomTraversal {
					// Perform all reads in random order.
					// indexes := rand.Perm(total)
				} else {
					// Perform all reads in sequential order.
				}
				wg.Done()
			}()
		}
		wg.Wait()
	}
}
<commit_msg>fix basic parallel benchmark<commit_after>package bolt

import (
	"errors"
	"fmt"
	"os"
	"sync"
	"testing"
)

const (
	BenchReadMode            = "read"
	BenchWriteMode           = "write"
	BenchSequentialTraversal = "sequential"
	BenchRandomTraversal     = "random"
)

type Benchmark struct {
	InputPath        string
	ReadWriteMode    string
	TraversalPattern string
	Parallelism      int
}

func NewBenchmark(inputPath, readWriteMode, traversalPattern string, parallelism int) *Benchmark {
	return &Benchmark{inputPath, readWriteMode, traversalPattern, parallelism}
}

func (bm *Benchmark) Run(b *testing.B) {

	// Open the database.
	db, err := Open(bm.InputPath, 0600)
	if err != nil {
		b.Fatalf("error: %+v", err)
		return
	}
	defer db.Close()

	buckets, err := buckets(db, bm.InputPath)
	if err != nil {
		b.Fatalf("error: %+v", err)
	}

	b.ResetTimer()

	// Keep running a fixed number of parallel reads until we run out of time.
	for i := 0; i < b.N; i++ {
		var wg sync.WaitGroup
		for j := 0; j < bm.Parallelism; j++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				if err := bm.runBuckets(b, db, buckets); err != nil {
					b.Fatalf("error: %+v", err)
				}
			}()
		}
		wg.Wait()
	}
}

// Run benchmark(s) for each of the given buckets.
func (bm *Benchmark) runBuckets(b *testing.B, db *DB, buckets []string) error {
	return db.View(func(tx *Tx) error {
		bucketsCount := len(buckets)
		for _, bucket := range buckets {
			c := tx.Bucket([]byte(bucket)).Cursor()
			count := 0
			for k, _ := c.First(); k != nil; k, _ = c.Next() {
				count++
			}
			if count != bucketsCount {
				return errors.New(fmt.Sprintf("wrong count: %d; expected: %d", count, bucketsCount))
			}
		}
		return nil
	})
}

func buckets(db *DB, path string) ([]string, error) {
	if _, err := os.Stat(path); os.IsNotExist(err) {
		return nil, err
	}

	buckets := []string{}

	err := db.View(func(tx *Tx) error {
		// Iterate over each bucket.
		return tx.ForEach(func(name []byte, _ *Bucket) error {
			buckets = append(buckets, string(name))
			return nil
		})
	})

	return buckets, err
}
<|endoftext|>


<commit_before>/*
Package bench provides a generic framework for performing latency benchmarks.
*/
package bench

import (
	"fmt"
	"os"
	"time"

	"github.com/codahale/hdrhistogram"
	"github.com/tsenart/tb"
)

const (
	maxRecordableLatencyNS = 300000000000
	sigFigs                = 5
)

// Requester synchronously issues requests for a particular system under test.
type Requester interface {
	// Setup prepares the Requester for benchmarking.
	Setup() error

	// Request performs a synchronous request to the system under test.
	Request() error

	// Teardown is called upon benchmark completion.
	Teardown() error
}

// Benchmark performs a system benchmark by issuing a certain number of
// requests at a specified rate and capturing the latency distribution.
type Benchmark struct {
	requester            Requester
	rateLimit            int64
	interval             time.Duration
	tb                   *tb.Bucket
	numRequests          int64
	histogram            *hdrhistogram.Histogram
	uncorrectedHistogram *hdrhistogram.Histogram
}

// NewBenchmark creates a Benchmark which runs a system benchmark using the
// given Requester. The rateLimit argument specifies the number of requests per
// interval to issue. A negative value disables rate limiting entirely. The
// numRequests argument specifies the total number of requests to issue in the
// benchmark.
func NewBenchmark(requester Requester, rateLimit int64, interval time.Duration,
	numRequests int64) *Benchmark {

	if interval <= 0 {
		interval = time.Second
	}

	return &Benchmark{
		requester:            requester,
		rateLimit:            rateLimit,
		interval:             interval,
		tb:                   tb.NewBucket(rateLimit, interval),
		numRequests:          numRequests,
		histogram:            hdrhistogram.New(0, maxRecordableLatencyNS, sigFigs),
		uncorrectedHistogram: hdrhistogram.New(0, maxRecordableLatencyNS, sigFigs),
	}
}

// Run the benchmark and return an error if something went wrong along the way.
// This can be called multiple times, overwriting the results on each call,
// unless Dispose is called.
func (b *Benchmark) Run() error {
	b.histogram.Reset()
	b.tb.Put(b.rateLimit)
	b.uncorrectedHistogram.Reset()

	if err := b.requester.Setup(); err != nil {
		return err
	}

	var err error
	if b.rateLimit < 0 {
		err = b.runFullThrottle()
	} else {
		err = b.runRateLimited()
	}

	if e := b.requester.Teardown(); e != nil {
		err = e
	}

	return err
}

// Dispose resources used by the Benchmark. Once this is called, this Benchmark
// can no longer be used.
func (b *Benchmark) Dispose() {
	b.tb.Close()
}

// GenerateLatencyDistribution generates a text file containing the specified
// latency distribution in a format plottable by
// http://hdrhistogram.github.io/HdrHistogram/plotFiles.html. If percentiles is
// nil, it defaults to a logarithmic percentile scale. If a rate limit was
// specified for the benchmark, this will also generate an uncorrected
// distribution file which does not account for coordinated omission.
func (b *Benchmark) GenerateLatencyDistribution(percentiles []float64, file string) error {
	if percentiles == nil {
		percentiles = defaultPercentiles
	}
	f, err := os.Create(file)
	if err != nil {
		return err
	}
	defer f.Close()

	f.WriteString("Value    Percentile    TotalCount    1/(1-Percentile)\n\n")
	for _, percentile := range percentiles {
		value := b.histogram.ValueAtQuantile(percentile)
		_, err := f.WriteString(fmt.Sprintf("%d    %f        %d            %f\n",
			value, percentile/100, 0, 1/(1-(percentile/100))))
		if err != nil {
			return err
		}
	}

	// Generate uncorrected distribution.
	if b.rateLimit > 0 {
		f, err := os.Create("uncorrected_" + file)
		if err != nil {
			return err
		}
		defer f.Close()

		f.WriteString("Value    Percentile    TotalCount    1/(1-Percentile)\n\n")
		for _, percentile := range percentiles {
			value := b.uncorrectedHistogram.ValueAtQuantile(percentile)
			_, err := f.WriteString(fmt.Sprintf("%d    %f        %d            %f\n",
				value, percentile/100, 0, 1/(1-(percentile/100))))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (b *Benchmark) runRateLimited() error {
	for i := int64(0); i < b.numRequests; i++ {
		b.tb.Wait(1)
		before := time.Now()
		if err := b.requester.Request(); err != nil {
			return err
		}
		latency := time.Since(before).Nanoseconds()
		if err := b.histogram.RecordCorrectedValue(latency, b.interval.Nanoseconds()); err != nil {
			return err
		}
		if err := b.uncorrectedHistogram.RecordValue(latency); err != nil {
			return err
		}
	}
	return nil
}

func (b *Benchmark) runFullThrottle() error {
	for i := int64(0); i < b.numRequests; i++ {
		before := time.Now()
		if err := b.requester.Request(); err != nil {
			return err
		}
		if err := b.histogram.RecordValue(time.Since(before).Nanoseconds()); err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Code cleanup<commit_after>/*
Package bench provides a generic framework for performing latency benchmarks.
*/
package bench

import (
	"fmt"
	"os"
	"time"

	"github.com/codahale/hdrhistogram"
	"github.com/tsenart/tb"
)

const (
	maxRecordableLatencyNS = 300000000000
	sigFigs                = 5
)

// Requester synchronously issues requests for a particular system under test.
type Requester interface {
	// Setup prepares the Requester for benchmarking.
	Setup() error

	// Request performs a synchronous request to the system under test.
	Request() error

	// Teardown is called upon benchmark completion.
	Teardown() error
}

// Benchmark performs a system benchmark by issuing a certain number of
// requests at a specified rate and capturing the latency distribution.
type Benchmark struct {
	requester            Requester
	rateLimit            int64
	interval             time.Duration
	tb                   *tb.Bucket
	numRequests          int64
	histogram            *hdrhistogram.Histogram
	uncorrectedHistogram *hdrhistogram.Histogram
}

// NewBenchmark creates a Benchmark which runs a system benchmark using the
// given Requester. The rateLimit argument specifies the number of requests per
// interval to issue. A negative value disables rate limiting entirely. The
// numRequests argument specifies the total number of requests to issue in the
// benchmark.
func NewBenchmark(requester Requester, rateLimit int64, interval time.Duration,
	numRequests int64) *Benchmark {

	if interval <= 0 {
		interval = time.Second
	}

	return &Benchmark{
		requester:            requester,
		rateLimit:            rateLimit,
		interval:             interval,
		tb:                   tb.NewBucket(rateLimit, interval),
		numRequests:          numRequests,
		histogram:            hdrhistogram.New(0, maxRecordableLatencyNS, sigFigs),
		uncorrectedHistogram: hdrhistogram.New(0, maxRecordableLatencyNS, sigFigs),
	}
}

// Run the benchmark and return an error if something went wrong along the way.
// This can be called multiple times, overwriting the results on each call,
// unless Dispose is called.
func (b *Benchmark) Run() error {
	b.histogram.Reset()
	b.uncorrectedHistogram.Reset()
	b.tb.Put(b.rateLimit)

	if err := b.requester.Setup(); err != nil {
		return err
	}

	var err error
	if b.rateLimit < 0 {
		err = b.runFullThrottle()
	} else {
		err = b.runRateLimited()
	}

	if e := b.requester.Teardown(); e != nil {
		err = e
	}

	return err
}

// Dispose resources used by the Benchmark. Once this is called, this Benchmark
// can no longer be used.
func (b *Benchmark) Dispose() {
	b.tb.Close()
}

// GenerateLatencyDistribution generates a text file containing the specified
// latency distribution in a format plottable by
// http://hdrhistogram.github.io/HdrHistogram/plotFiles.html. If percentiles is
// nil, it defaults to a logarithmic percentile scale. If a rate limit was
// specified for the benchmark, this will also generate an uncorrected
// distribution file which does not account for coordinated omission.
func (b *Benchmark) GenerateLatencyDistribution(percentiles []float64, file string) error {
	if percentiles == nil {
		percentiles = defaultPercentiles
	}
	f, err := os.Create(file)
	if err != nil {
		return err
	}
	defer f.Close()

	f.WriteString("Value    Percentile    TotalCount    1/(1-Percentile)\n\n")
	for _, percentile := range percentiles {
		value := b.histogram.ValueAtQuantile(percentile)
		_, err := f.WriteString(fmt.Sprintf("%d    %f        %d            %f\n",
			value, percentile/100, 0, 1/(1-(percentile/100))))
		if err != nil {
			return err
		}
	}

	// Generate uncorrected distribution.
	if b.rateLimit > 0 {
		f, err := os.Create("uncorrected_" + file)
		if err != nil {
			return err
		}
		defer f.Close()

		f.WriteString("Value    Percentile    TotalCount    1/(1-Percentile)\n\n")
		for _, percentile := range percentiles {
			value := b.uncorrectedHistogram.ValueAtQuantile(percentile)
			_, err := f.WriteString(fmt.Sprintf("%d    %f        %d            %f\n",
				value, percentile/100, 0, 1/(1-(percentile/100))))
			if err != nil {
				return err
			}
		}
	}

	return nil
}

func (b *Benchmark) runRateLimited() error {
	for i := int64(0); i < b.numRequests; i++ {
		b.tb.Wait(1)
		before := time.Now()
		if err := b.requester.Request(); err != nil {
			return err
		}
		latency := time.Since(before).Nanoseconds()
		if err := b.histogram.RecordCorrectedValue(latency, b.interval.Nanoseconds()); err != nil {
			return err
		}
		if err := b.uncorrectedHistogram.RecordValue(latency); err != nil {
			return err
		}
	}
	return nil
}

func (b *Benchmark) runFullThrottle() error {
	for i := int64(0); i < b.numRequests; i++ {
		before := time.Now()
		if err := b.requester.Request(); err != nil {
			return err
		}
		if err := b.histogram.RecordValue(time.Since(before).Nanoseconds()); err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>