<commit_before>package resource

import (
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"time"

	msgpack "gopkg.in/vmihailenco/msgpack.v2"

	"log"

	"github.com/gbl08ma/sqalx"
	"github.com/underlx/disturbancesmlx/dataobjects"
	"github.com/yarf-framework/yarf"
)

type resource struct {
	yarf.Resource
	node    sqalx.Node
	hashKey []byte
}

// Beginx is shorthand for resource.node.Beginx()
func (r *resource) Beginx() (sqalx.Node, error) {
	return r.node.Beginx()
}

// DecodeRequest decodes a request according to its headers and places the result in v
func (r *resource) DecodeRequest(c *yarf.Context, v interface{}) error {
	contentType := c.Request.Header.Get("Content-Type")
	var err error
	switch {
	case strings.Contains(contentType, "msgpack"):
		err = msgpack.NewDecoder(c.Request.Body).Decode(v)
	case strings.Contains(contentType, "json"):
		fallthrough
	default:
		err = json.NewDecoder(c.Request.Body).Decode(v)
	}

	if err != nil {
		return &yarf.CustomError{
			HTTPCode:  http.StatusBadRequest,
			ErrorMsg:  "Failed to decode request",
			ErrorBody: err.Error(),
		}
	}
	return nil
}

// AuthenticateClient authenticates a API client
func (r *resource) AuthenticateClient(c *yarf.Context) (pair *dataobjects.APIPair, err error) {
	tx, err := r.node.Beginx()
	if err != nil {
		return nil, err
	}
	defer tx.Commit() // read-only tx

	key, secret, ok := c.Request.BasicAuth()
	if !ok {
		return nil, errors.New("Missing authorization header")
	}

	pair, err = dataobjects.GetPairIfCorrect(tx, key, secret, r.hashKey)
	if err != nil {
		return nil, errors.New("Incorrect authorization")
	}

	return pair, nil
}

// RenderUnauthorized writes a 401 unauthorized to the response
// and requests authentication
func RenderUnauthorized(c *yarf.Context) {
	c.Response.Header().Set("WWW-Authenticate", `Basic realm="api"`)
	c.Response.WriteHeader(401)
	c.Response.Write([]byte("Unauthorized.\n"))
}

// RenderData takes a interface{} object and writes the encoded representation of it.
// Encoding used will be idented JSON, non-idented JSON, Msgpack or XML
func RenderData(c *yarf.Context, data interface{}, cacheControl string) {
	accept := c.Request.Header.Get("Accept")
	switch {
	case strings.Contains(accept, "json"):
		c.Response.Header().Set("Content-Type", "application/json; charset=utf-8")
		c.RenderJSON(data)
	case strings.Contains(accept, "xml") && !strings.Contains(accept, "xhtml"):
		c.Response.Header().Set("Content-Type", "application/xml; charset=utf-8")
		c.RenderXML(data)
	case strings.Contains(accept, "msgpack"):
		RenderMsgpack(c, data)
	default:
		c.Response.Header().Set("Content-Type", "application/json; charset=utf-8")
		c.RenderJSONIndent(data)
	}
	if cacheControl != "" {
		c.Response.Header().Set("Cache-Control", cacheControl)
	}
}

// RenderMsgpack takes a interface{} object and writes the Msgpack encoded string of it.
func RenderMsgpack(c *yarf.Context, data interface{}) {
	c.Response.Header().Set("Content-Type", "application/msgpack")
	// Set content
	encoded, err := msgpack.Marshal(data)
	if err != nil {
		log.Println(err)
		c.Response.Write([]byte(err.Error()))
	} else {
		c.Response.Write(encoded)
	}
}

func maxDuration(d1 time.Duration, d2 time.Duration) time.Duration {
	if d1 > d2 {
		return d1
	}
	return d2
}
<commit_msg>API: Fix header sending<commit_after>package resource

import (
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"time"

	msgpack "gopkg.in/vmihailenco/msgpack.v2"

	"log"

	"github.com/gbl08ma/sqalx"
	"github.com/underlx/disturbancesmlx/dataobjects"
	"github.com/yarf-framework/yarf"
)

type resource struct {
	yarf.Resource
	node    sqalx.Node
	hashKey []byte
}

// Beginx is shorthand for resource.node.Beginx()
func (r *resource) Beginx() (sqalx.Node, error) {
	return r.node.Beginx()
}

// DecodeRequest decodes a request according to its headers and places the result in v
func (r *resource) DecodeRequest(c *yarf.Context, v interface{}) error {
	contentType := c.Request.Header.Get("Content-Type")
	var err error
	switch {
	case strings.Contains(contentType, "msgpack"):
		err = msgpack.NewDecoder(c.Request.Body).Decode(v)
	case strings.Contains(contentType, "json"):
		fallthrough
	default:
		err = json.NewDecoder(c.Request.Body).Decode(v)
	}

	if err != nil {
		return &yarf.CustomError{
			HTTPCode:  http.StatusBadRequest,
			ErrorMsg:  "Failed to decode request",
			ErrorBody: err.Error(),
		}
	}
	return nil
}

// AuthenticateClient authenticates a API client
func (r *resource) AuthenticateClient(c *yarf.Context) (pair *dataobjects.APIPair, err error) {
	tx, err := r.node.Beginx()
	if err != nil {
		return nil, err
	}
	defer tx.Commit() // read-only tx

	key, secret, ok := c.Request.BasicAuth()
	if !ok {
		return nil, errors.New("Missing authorization header")
	}

	pair, err = dataobjects.GetPairIfCorrect(tx, key, secret, r.hashKey)
	if err != nil {
		return nil, errors.New("Incorrect authorization")
	}

	return pair, nil
}

// RenderUnauthorized writes a 401 unauthorized to the response
// and requests authentication
func RenderUnauthorized(c *yarf.Context) {
	c.Response.Header().Set("WWW-Authenticate", `Basic realm="api"`)
	c.Response.WriteHeader(401)
	c.Response.Write([]byte("Unauthorized.\n"))
}

// RenderData takes a interface{} object and writes the encoded representation of it.
// Encoding used will be idented JSON, non-idented JSON, Msgpack or XML
func RenderData(c *yarf.Context, data interface{}, cacheControl string) {
	if cacheControl != "" {
		c.Response.Header().Set("Cache-Control", cacheControl)
	}
	accept := c.Request.Header.Get("Accept")
	switch {
	case strings.Contains(accept, "json"):
		c.Response.Header().Set("Content-Type", "application/json; charset=utf-8")
		c.RenderJSON(data)
	case strings.Contains(accept, "xml") && !strings.Contains(accept, "xhtml"):
		c.Response.Header().Set("Content-Type", "application/xml; charset=utf-8")
		c.RenderXML(data)
	case strings.Contains(accept, "msgpack"):
		RenderMsgpack(c, data)
	default:
		c.Response.Header().Set("Content-Type", "application/json; charset=utf-8")
		c.RenderJSONIndent(data)
	}
}

// RenderMsgpack takes a interface{} object and writes the Msgpack encoded string of it.
func RenderMsgpack(c *yarf.Context, data interface{}) {
	c.Response.Header().Set("Content-Type", "application/msgpack")
	// Set content
	encoded, err := msgpack.Marshal(data)
	if err != nil {
		log.Println(err)
		c.Response.Write([]byte(err.Error()))
	} else {
		c.Response.Write(encoded)
	}
}

func maxDuration(d1 time.Duration, d2 time.Duration) time.Duration {
	if d1 > d2 {
		return d1
	}
	return d2
}
<|endoftext|>


<commit_before>package bloom

import (
	"hash/fnv"
)

type BloomFilter struct {
	values   map[uint64]bool
	capacity uint64
}

func str_to_fnv(s string) (uint64, uint64) {
	h1 := fnv.New64()
	h1.Write([]byte(s))

	h2 := fnv.New64a()
	h2.Write([]byte(s))

	return h1.Sum64(), h2.Sum64()
}

func New(capacity int) *BloomFilter {
	return &BloomFilter{make(map[uint64]bool, capacity), uint64(capacity)}
}

func (bf *BloomFilter) Add(s string) {
	h1, h2 := str_to_fnv(s)
	bf.values[h1%bf.capacity] = true
	bf.values[h2%bf.capacity] = true
}

func (bf *BloomFilter) IsPresent(s string) bool {
	h1, h2 := str_to_fnv(s)

	return bf.values[h1%bf.capacity] && bf.values[h2%bf.capacity]
}
<commit_msg>Extract get() and set()<commit_after>package bloom

import (
	"hash/fnv"
)

type BloomFilter struct {
	values   map[uint64]bool
	capacity uint64
}

func str_to_fnv(s string) (uint64, uint64) {
	h1 := fnv.New64()
	h1.Write([]byte(s))

	h2 := fnv.New64a()
	h2.Write([]byte(s))

	return h1.Sum64(), h2.Sum64()
}

func (bf *BloomFilter) set(h uint64, value bool) {
	bf.values[h%bf.capacity] = value
}

func (bf *BloomFilter) get(h uint64) bool {
	return bf.values[h%bf.capacity]
}

func New(capacity int) *BloomFilter {
	return &BloomFilter{make(map[uint64]bool, capacity), uint64(capacity)}
}

func (bf *BloomFilter) Add(s string) {
	h1, h2 := str_to_fnv(s)
	bf.set(h1, true)
	bf.set(h2, true)
}

func (bf *BloomFilter) IsPresent(s string) bool {
	h1, h2 := str_to_fnv(s)

	return bf.get(h1) && bf.get(h2)
}
<|endoftext|>


<commit_before>// Throttler fills the gap between sync.WaitGroup and manually monitoring your goroutines
// with channels. The API is almost identical to Wait Groups, but it allows you to set
// a max number of workers that can be running simultaneously. It uses channels internally
// to block until a job completes by calling Done() or until all jobs have been completed.
//
// Compare the Throttler example to the sync.WaitGroup example http://golang.org/pkg/sync/#example_WaitGroup
//
// See a fully functional example on the playground at http://bit.ly/throttler-docs
package throttler

type Throttler struct {
	maxWorkers    int
	workerCount   int
	totalJobs     int
	jobsStarted   int
	jobsCompleted int
	doneChan      chan bool
}

// New returns a Throttler that will govern the max number of workers and will
// work with the total number of jobs. It panics if maxWorkers < 1.
func New(maxWorkers, totalJobs int) *Throttler {
	if maxWorkers < 1 {
		panic("maxWorkers has to be at least 1")
	}
	return &Throttler{
		maxWorkers: maxWorkers,
		totalJobs:  totalJobs,
		doneChan:   make(chan bool, totalJobs),
	}
}

// Throttle works similarly to sync.WaitGroup, except inside your goroutine dispatch
// loop rather than after. It will not block until the number of active workers
// matches the max number of workers designated in the call to NewThrottler or
// all of the jobs have been dispatched. It stops blocking when Done has been called
// as many times as totalJobs.
func (t *Throttler) Throttle() {
	if t.totalJobs < 1 {
		return
	}
	t.jobsStarted++
	t.workerCount++

	if t.workerCount == t.maxWorkers {
		<-t.doneChan
		t.jobsCompleted++
		t.workerCount--
	}

	if t.jobsStarted == t.totalJobs {
		for t.jobsCompleted < t.totalJobs {
			<-t.doneChan
			t.jobsCompleted++
		}
	}
}

// Done lets Throttler know that a job has been completed so that another worker
// can be activated. If Done is called less times than totalJobs,
// Throttle will block forever
func (t *Throttler) Done() {
	t.doneChan <- true
}
<commit_msg>Replaced doneChan chan bool with chan struct{}<commit_after>// Throttler fills the gap between sync.WaitGroup and manually monitoring your goroutines
// with channels. The API is almost identical to Wait Groups, but it allows you to set
// a max number of workers that can be running simultaneously. It uses channels internally
// to block until a job completes by calling Done() or until all jobs have been completed.
//
// Compare the Throttler example to the sync.WaitGroup example http://golang.org/pkg/sync/#example_WaitGroup
//
// See a fully functional example on the playground at http://bit.ly/throttler-docs
package throttler

type Throttler struct {
	maxWorkers    int
	workerCount   int
	totalJobs     int
	jobsStarted   int
	jobsCompleted int
	doneChan      chan struct{}
}

// New returns a Throttler that will govern the max number of workers and will
// work with the total number of jobs. It panics if maxWorkers < 1.
func New(maxWorkers, totalJobs int) *Throttler {
	if maxWorkers < 1 {
		panic("maxWorkers has to be at least 1")
	}
	return &Throttler{
		maxWorkers: maxWorkers,
		totalJobs:  totalJobs,
		doneChan:   make(chan struct{}, totalJobs),
	}
}

// Throttle works similarly to sync.WaitGroup, except inside your goroutine dispatch
// loop rather than after. It will not block until the number of active workers
// matches the max number of workers designated in the call to NewThrottler or
// all of the jobs have been dispatched. It stops blocking when Done has been called
// as many times as totalJobs.
func (t *Throttler) Throttle() {
	if t.totalJobs < 1 {
		return
	}
	t.jobsStarted++
	t.workerCount++

	if t.workerCount == t.maxWorkers {
		<-t.doneChan
		t.jobsCompleted++
		t.workerCount--
	}

	if t.jobsStarted == t.totalJobs {
		for t.jobsCompleted < t.totalJobs {
			<-t.doneChan
			t.jobsCompleted++
		}
	}
}

// Done lets Throttler know that a job has been completed so that another worker
// can be activated. If Done is called less times than totalJobs,
// Throttle will block forever
func (t *Throttler) Done() {
	t.doneChan <- struct{}{}
}
<|endoftext|>


<commit_before><commit_msg>Added rates and exp table<commit_after><|endoftext|>


<commit_before>package purge

import (
	"context"
	"math"
	"strconv"
	"time"

	"github.com/fsamin/go-dump"
	"github.com/go-gorp/gorp"
	"github.com/rockbears/log"
	"go.opencensus.io/stats"

	"github.com/ovh/cds/engine/api/application"
	"github.com/ovh/cds/engine/api/database/gorpmapping"
	"github.com/ovh/cds/engine/api/event"
	"github.com/ovh/cds/engine/api/repositoriesmanager"
	"github.com/ovh/cds/engine/api/workflow"
	"github.com/ovh/cds/engine/cache"
	"github.com/ovh/cds/engine/featureflipping"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/luascript"
	"github.com/ovh/cds/sdk/telemetry"
)

type MarkAsDeleteOptions struct {
	DryRun bool
}

const (
	RunStatus          = "run_status"
	RunDaysBefore      = "run_days_before"
	RunHasGitBranch    = "has_git_branch"
	RunGitBranchExist  = "git_branch_exist"
	RunChangeExist     = "gerrit_change_exist"
	RunChangeMerged    = "gerrit_change_merged"
	RunChangeAbandoned = "gerrit_change_abandoned"
	RunChangeDayBefore = "gerrit_change_days_before"
)

func GetRetentionPolicyVariables() []string {
	return []string{RunDaysBefore, RunStatus, RunHasGitBranch, RunGitBranchExist, RunChangeMerged, RunChangeAbandoned, RunChangeDayBefore, RunChangeExist}
}

func markWorkflowRunsToDelete(ctx context.Context, store cache.Store, db *gorp.DbMap, workflowRunsMarkToDelete *stats.Int64Measure) error {
	ctx, end := telemetry.Span(ctx, "purge.markWorkflowRunsToDelete")
	defer end()

	dao := new(workflow.WorkflowDAO)
	wfs, err := dao.LoadAll(ctx, db)
	if err != nil {
		return err
	}
	for _, wf := range wfs {
		_, enabled := featureflipping.IsEnabled(ctx, gorpmapping.Mapper, db, sdk.FeaturePurgeName, map[string]string{"project_key": wf.ProjectKey})
		if !enabled {
			continue
		}
		if err := ApplyRetentionPolicyOnWorkflow(ctx, store, db, wf, MarkAsDeleteOptions{DryRun: false}, nil); err != nil {
			ctx = sdk.ContextWithStacktrace(ctx, err)
			log.Error(ctx, "%v", err)
		}
	}
	workflow.CountWorkflowRunsMarkToDelete(ctx, db, workflowRunsMarkToDelete)
	return nil
}

func ApplyRetentionPolicyOnWorkflow(ctx context.Context, store cache.Store, db *gorp.DbMap, wf sdk.Workflow, opts MarkAsDeleteOptions, u *sdk.AuthentifiedUser) error {
	ctx, end := telemetry.Span(ctx, "purge.ApplyRetentionPolicyOnWorkflow")
	defer end()

	var vcsClient sdk.VCSAuthorizedClientService
	var app sdk.Application
	if wf.WorkflowData.Node.Context != nil {
		appID := wf.WorkflowData.Node.Context.ApplicationID
		if appID != 0 {
			appDB, err := application.LoadByID(db, appID)
			if err != nil {
				return err
			}
			app = *appDB
			if app.RepositoryFullname != "" {
				tx, err := db.Begin()
				if err != nil {
					return sdk.WithStack(err)
				}
				//Get the RepositoriesManager Client
				vcsServer, err := repositoriesmanager.LoadProjectVCSServerLinkByProjectKeyAndVCSServerName(ctx, tx, wf.ProjectKey, app.VCSServer)
				if err != nil {
					_ = tx.Rollback()
					return err
				}
				vcsClient, err = repositoriesmanager.AuthorizedClient(ctx, tx, store, wf.ProjectKey, vcsServer)
				if err != nil {
					_ = tx.Rollback()
					return sdk.WithStack(err)
				}
				if err := tx.Commit(); err != nil {
					_ = tx.Rollback()
					return err
				}
			}
		}
	}

	branchesMap := make(map[string]struct{})
	if vcsClient != nil {
		var err error
		branchesMap, err = getBranches(ctx, app.RepositoryFullname, vcsClient)
		if err != nil {
			return err
		}
	}

	runs := make([]sdk.WorkflowRunToKeep, 0)
	var nbRunsAnalyzed int64
	limit := 50
	offset := 0
	for {
		wfRuns, _, _, count, err := workflow.LoadRunsSummaries(ctx, db, wf.ProjectKey, wf.Name, offset, limit, nil)
		if err != nil {
			return err
		}

		nbRunsAnalyzed = int64(len(wfRuns))
		for _, run := range wfRuns {

			payload, err := extractPayload(run)
			if err != nil {
				return err
			}

			var forkBranches map[string]struct{}
			isFork := false
			if gitRepo, has := payload["git.repository"]; has {
				if gitRepo != app.RepositoryFullname {
					isFork = true
					forkBranches, err = getBranches(ctx, gitRepo, vcsClient)
					if err != nil {
						return err
					}
				}
			}

			var keep bool
			if !isFork {
				keep, err = applyRetentionPolicyOnRun(ctx, db, wf, run, payload, branchesMap, app, vcsClient, opts)
			} else {
				keep, err = applyRetentionPolicyOnRun(ctx, db, wf, run, payload, forkBranches, app, vcsClient, opts)
			}
			if err != nil {
				return err
			}
			if keep {
				runs = append(runs, sdk.WorkflowRunToKeep{ID: run.ID, Num: run.Number, Status: run.Status})
			}
		}

		if count > offset+limit {
			offset += limit
			if u != nil {
				event.PublishWorkflowRetentionDryRun(ctx, wf.ProjectKey, wf.Name, "INCOMING", "", runs, nbRunsAnalyzed, u)
				runs = runs[:0]
			}
			continue
		}
		break
	}
	if u != nil {
		event.PublishWorkflowRetentionDryRun(ctx, wf.ProjectKey, wf.Name, "DONE", "", runs, nbRunsAnalyzed, u)
	}
	return nil
}

func getBranches(ctx context.Context, repo string, vcsClient sdk.VCSAuthorizedClientService) (map[string]struct{}, error) {
	branchesMap := make(map[string]struct{})
	branches, err := vcsClient.Branches(ctx, repo, sdk.VCSBranchesFilter{})
	if err != nil {
		return nil, err
	}
	log.Info(ctx, "Purge getting branch for repo %s - count: %d", repo, len(branches))
	defaultBranchFound := false
	for _, b := range branches {
		branchesMap[b.DisplayID] = struct{}{}
		if b.Default {
			log.Info(ctx, "Getting default branch for repo %s - %s", repo, b.DisplayID)
			defaultBranchFound = true
		}
	}

	if !defaultBranchFound {
		log.Warn(ctx, "Getting default branch for repo %s - not found", repo)
	}
	return branchesMap, nil
}

func extractPayload(run sdk.WorkflowRunSummary) (map[string]string, error) {
	vars := make(map[string]string)
	// Add payload as variable
	if run.ToCraftOpts != nil {
		switch {
		case run.ToCraftOpts.Manual != nil:
			payload := run.ToCraftOpts.Manual.Payload
			if payload != nil {
				// COMPUTE PAYLOAD
				e := dump.NewDefaultEncoder()
				e.Formatters = []dump.KeyFormatterFunc{dump.WithDefaultLowerCaseFormatter()}
				e.ExtraFields.DetailedMap = false
				e.ExtraFields.DetailedStruct = false
				e.ExtraFields.Len = false
				e.ExtraFields.Type = false
				tmpVars, err := e.ToStringMap(payload)
				if err != nil {
					return nil, sdk.WithStack(err)
				}
				for k, v := range tmpVars {
					vars[k] = v
				}
			}
		case run.ToCraftOpts.Hook != nil && run.ToCraftOpts.Hook.Payload != nil:
			vars = run.ToCraftOpts.Hook.Payload
		}
	}
	return vars, nil
}

func applyRetentionPolicyOnRun(ctx context.Context, db *gorp.DbMap, wf sdk.Workflow, run sdk.WorkflowRunSummary, payload map[string]string, branchesMap map[string]struct{}, app sdk.Application, vcsClient sdk.VCSAuthorizedClientService, opts MarkAsDeleteOptions) (bool, error) {
	if wf.ToDelete && !opts.DryRun {
		if err := workflow.MarkWorkflowRunsAsDelete(db, []int64{run.ID}); err != nil {
			return true, sdk.WithStack(err)
		}
		return false, nil
	}

	luaCheck, err := luascript.NewCheck()
	if err != nil {
		return true, sdk.WithStack(err)
	}

	if err := purgeComputeVariables(ctx, luaCheck, run, payload, branchesMap, app, vcsClient); err != nil {
		return true, err
	}

	retentionPolicy := defaultRunRetentionPolicy
	if wf.RetentionPolicy != "" {
		retentionPolicy = wf.RetentionPolicy
	}

	// Enabling strict checks on variables to prevent errors on rule definition
	if err := luaCheck.EnableStrict(); err != nil {
		return true, sdk.WithStack(err)
	}

	if err := luaCheck.Perform(retentionPolicy); err != nil {
		return true, sdk.NewErrorFrom(sdk.ErrWrongRequest, "unable to apply retention policy on workflow %s/%s: %v", wf.ProjectKey, wf.Name, err)
	}

	if luaCheck.Result {
		return true, nil
	}
	if !opts.DryRun {
		if err := workflow.MarkWorkflowRunsAsDelete(db, []int64{run.ID}); err != nil {
			return true, sdk.WithStack(err)
		}
	}
	return false, nil
}

func purgeComputeVariables(ctx context.Context, luaCheck *luascript.Check, run sdk.WorkflowRunSummary, payload map[string]string, branchesMap map[string]struct{}, app sdk.Application, vcsClient sdk.VCSAuthorizedClientService) error {
	vars := payload
	varsFloats := make(map[string]float64)

	// If we have gerrit change id, check status
	if changeID, ok := vars["gerrit.change.id"]; ok && vcsClient != nil {
		ch, err := vcsClient.PullRequest(ctx, app.RepositoryFullname, changeID)
		if err != nil {
			if !sdk.ErrorIs(err, sdk.ErrNotFound) {
				return err
			}
			vars[RunChangeExist] = "false"
		} else {
			vars[RunChangeExist] = "true"
			vars[RunChangeMerged] = strconv.FormatBool(ch.Merged)
			vars[RunChangeAbandoned] = strconv.FormatBool(ch.Closed)
			varsFloats[RunChangeDayBefore] = math.Floor(time.Now().Sub(ch.Updated).Hours())
		}
	}

	// If we have a branch in payload, check if it exists on repository branches list
	b, has := vars["git.branch"]
	var exist bool
	if has {
		_, exist = branchesMap[b]
	}
	vars[RunHasGitBranch] = strconv.FormatBool(has)
	vars[RunGitBranchExist] = strconv.FormatBool(exist)

	vars[RunStatus] = run.Status

	varsFloats[RunDaysBefore] = math.Floor(time.Now().Sub(run.LastModified).Hours() / 24)

	luaCheck.SetVariables(vars)
	luaCheck.SetFloatVariables(varsFloats)
	return nil
}
<commit_msg>fix(api): purge check vcsclient (#6019)<commit_after>package purge

import (
	"context"
	"math"
	"strconv"
	"time"

	"github.com/fsamin/go-dump"
	"github.com/go-gorp/gorp"
	"github.com/rockbears/log"
	"go.opencensus.io/stats"

	"github.com/ovh/cds/engine/api/application"
	"github.com/ovh/cds/engine/api/database/gorpmapping"
	"github.com/ovh/cds/engine/api/event"
	"github.com/ovh/cds/engine/api/repositoriesmanager"
	"github.com/ovh/cds/engine/api/workflow"
	"github.com/ovh/cds/engine/cache"
	"github.com/ovh/cds/engine/featureflipping"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/luascript"
	"github.com/ovh/cds/sdk/telemetry"
)

type MarkAsDeleteOptions struct {
	DryRun bool
}

const (
	RunStatus          = "run_status"
	RunDaysBefore      = "run_days_before"
	RunHasGitBranch    = "has_git_branch"
	RunGitBranchExist  = "git_branch_exist"
	RunChangeExist     = "gerrit_change_exist"
	RunChangeMerged    = "gerrit_change_merged"
	RunChangeAbandoned = "gerrit_change_abandoned"
	RunChangeDayBefore = "gerrit_change_days_before"
)

func GetRetentionPolicyVariables() []string {
	return []string{RunDaysBefore, RunStatus, RunHasGitBranch, RunGitBranchExist, RunChangeMerged, RunChangeAbandoned, RunChangeDayBefore, RunChangeExist}
}

func markWorkflowRunsToDelete(ctx context.Context, store cache.Store, db *gorp.DbMap, workflowRunsMarkToDelete *stats.Int64Measure) error {
	ctx, end := telemetry.Span(ctx, "purge.markWorkflowRunsToDelete")
	defer end()

	dao := new(workflow.WorkflowDAO)
	wfs, err := dao.LoadAll(ctx, db)
	if err != nil {
		return err
	}
	for _, wf := range wfs {
		_, enabled := featureflipping.IsEnabled(ctx, gorpmapping.Mapper, db, sdk.FeaturePurgeName, map[string]string{"project_key": wf.ProjectKey})
		if !enabled {
			continue
		}
		if err := ApplyRetentionPolicyOnWorkflow(ctx, store, db, wf, MarkAsDeleteOptions{DryRun: false}, nil); err != nil {
			ctx = sdk.ContextWithStacktrace(ctx, err)
			log.Error(ctx, "%v", err)
		}
	}
	workflow.CountWorkflowRunsMarkToDelete(ctx, db, workflowRunsMarkToDelete)
	return nil
}

func ApplyRetentionPolicyOnWorkflow(ctx context.Context, store cache.Store, db *gorp.DbMap, wf sdk.Workflow, opts MarkAsDeleteOptions, u *sdk.AuthentifiedUser) error {
	ctx, end := telemetry.Span(ctx, "purge.ApplyRetentionPolicyOnWorkflow")
	defer end()

	var vcsClient sdk.VCSAuthorizedClientService
	var app sdk.Application
	if wf.WorkflowData.Node.Context != nil {
		appID := wf.WorkflowData.Node.Context.ApplicationID
		if appID != 0 {
			appDB, err := application.LoadByID(db, appID)
			if err != nil {
				return err
			}
			app = *appDB
			if app.RepositoryFullname != "" {
				tx, err := db.Begin()
				if err != nil {
					return sdk.WithStack(err)
				}
				//Get the RepositoriesManager Client
				vcsServer, err := repositoriesmanager.LoadProjectVCSServerLinkByProjectKeyAndVCSServerName(ctx, tx, wf.ProjectKey, app.VCSServer)
				if err != nil {
					_ = tx.Rollback()
					return err
				}
				vcsClient, err = repositoriesmanager.AuthorizedClient(ctx, tx, store, wf.ProjectKey, vcsServer)
				if err != nil {
					_ = tx.Rollback()
					return sdk.WithStack(err)
				}
				if err := tx.Commit(); err != nil {
					_ = tx.Rollback()
					return err
				}
			}
		}
	}

	branchesMap := make(map[string]struct{})
	if vcsClient != nil {
		var err error
		branchesMap, err = getBranches(ctx, app.RepositoryFullname, vcsClient)
		if err != nil {
			return err
		}
	}

	runs := make([]sdk.WorkflowRunToKeep, 0)
	var nbRunsAnalyzed int64
	limit := 50
	offset := 0
	for {
		wfRuns, _, _, count, err := workflow.LoadRunsSummaries(ctx, db, wf.ProjectKey, wf.Name, offset, limit, nil)
		if err != nil {
			return err
		}

		nbRunsAnalyzed = int64(len(wfRuns))
		for _, run := range wfRuns {

			payload, err := extractPayload(run)
			if err != nil {
				return err
			}

			var forkBranches map[string]struct{}
			isFork := false
			if gitRepo, has := payload["git.repository"]; has {
				if gitRepo != app.RepositoryFullname {
					isFork = true
					if vcsClient != nil {
						forkBranches, err = getBranches(ctx, gitRepo, vcsClient)
						if err != nil {
							return err
						}
					}
				}
			}

			var keep bool
			if isFork {
				keep, err = applyRetentionPolicyOnRun(ctx, db, wf, run, payload, forkBranches, app, vcsClient, opts)
			} else {
				keep, err = applyRetentionPolicyOnRun(ctx, db, wf, run, payload, branchesMap, app, vcsClient, opts)
			}
			if keep {
				runs = append(runs, sdk.WorkflowRunToKeep{ID: run.ID, Num: run.Number, Status: run.Status})
			}
			if err != nil {
				log.Error(ctx, "error on run %v:%d err:%v", wf.Name, run.Number, err)
				continue

			}
		}

		if count > offset+limit {
			offset += limit
			if u != nil {
				event.PublishWorkflowRetentionDryRun(ctx, wf.ProjectKey, wf.Name, "INCOMING", "", runs, nbRunsAnalyzed, u)
				runs = runs[:0]
			}
			continue
		}
		break
	}
	if u != nil {
		event.PublishWorkflowRetentionDryRun(ctx, wf.ProjectKey, wf.Name, "DONE", "", runs, nbRunsAnalyzed, u)
	}
	return nil
}

func getBranches(ctx context.Context, repo string, vcsClient sdk.VCSAuthorizedClientService) (map[string]struct{}, error) {
	branchesMap := make(map[string]struct{})
	branches, err := vcsClient.Branches(ctx, repo, sdk.VCSBranchesFilter{})
	if err != nil {
		return nil, err
	}
	log.Info(ctx, "Purge getting branch for repo %s - count: %d", repo, len(branches))
	defaultBranchFound := false
	for _, b := range branches {
		branchesMap[b.DisplayID] = struct{}{}
		if b.Default {
			log.Info(ctx, "Getting default branch for repo %s - %s", repo, b.DisplayID)
			defaultBranchFound = true
		}
	}

	if !defaultBranchFound {
		log.Warn(ctx, "Getting default branch for repo %s - not found", repo)
	}
	return branchesMap, nil
}

func extractPayload(run sdk.WorkflowRunSummary) (map[string]string, error) {
	vars := make(map[string]string)
	// Add payload as variable
	if run.ToCraftOpts != nil {
		switch {
		case run.ToCraftOpts.Manual != nil:
			payload := run.ToCraftOpts.Manual.Payload
			if payload != nil {
				// COMPUTE PAYLOAD
				e := dump.NewDefaultEncoder()
				e.Formatters = []dump.KeyFormatterFunc{dump.WithDefaultLowerCaseFormatter()}
				e.ExtraFields.DetailedMap = false
				e.ExtraFields.DetailedStruct = false
				e.ExtraFields.Len = false
				e.ExtraFields.Type = false
				tmpVars, err := e.ToStringMap(payload)
				if err != nil {
					return nil, sdk.WithStack(err)
				}
				for k, v := range tmpVars {
					vars[k] = v
				}
			}
		case run.ToCraftOpts.Hook != nil && run.ToCraftOpts.Hook.Payload != nil:
			vars = run.ToCraftOpts.Hook.Payload
		}
	}
	return vars, nil
}

func applyRetentionPolicyOnRun(ctx context.Context, db *gorp.DbMap, wf sdk.Workflow, run sdk.WorkflowRunSummary, payload map[string]string, branchesMap map[string]struct{}, app sdk.Application, vcsClient sdk.VCSAuthorizedClientService, opts MarkAsDeleteOptions) (bool, error) {
	if wf.ToDelete && !opts.DryRun {
		if err := workflow.MarkWorkflowRunsAsDelete(db, []int64{run.ID}); err != nil {
			return true, sdk.WithStack(err)
		}
		return false, nil
	}

	luaCheck, err := luascript.NewCheck()
	if err != nil {
		return true, sdk.WithStack(err)
	}

	if err := purgeComputeVariables(ctx, luaCheck, run, payload, branchesMap, app, vcsClient); err != nil {
		return true, err
	}

	retentionPolicy := defaultRunRetentionPolicy
	if wf.RetentionPolicy != "" {
		retentionPolicy = wf.RetentionPolicy
	}

	// Enabling strict checks on variables to prevent errors on rule definition
	if err := luaCheck.EnableStrict(); err != nil {
		return true, sdk.WithStack(err)
	}

	if err := luaCheck.Perform(retentionPolicy); err != nil {
		return true, sdk.NewErrorFrom(sdk.ErrWrongRequest, "unable to apply retention policy on workflow %s/%s: %v", wf.ProjectKey, wf.Name, err)
	}

	if luaCheck.Result {
		return true, nil
	}
	if !opts.DryRun {
		if err := workflow.MarkWorkflowRunsAsDelete(db, []int64{run.ID}); err != nil {
			return true, sdk.WithStack(err)
		}
	}
	return false, nil
}

func purgeComputeVariables(ctx context.Context, luaCheck *luascript.Check, run sdk.WorkflowRunSummary, payload map[string]string, branchesMap map[string]struct{}, app sdk.Application, vcsClient sdk.VCSAuthorizedClientService) error {
	vars := payload
	varsFloats := make(map[string]float64)

	// If we have gerrit change id, check status
	if changeID, ok := vars["gerrit.change.id"]; ok && vcsClient != nil {
		ch, err := vcsClient.PullRequest(ctx, app.RepositoryFullname, changeID)
		if err != nil {
			if !sdk.ErrorIs(err, sdk.ErrNotFound) {
				return err
			}
			vars[RunChangeExist] = "false"
		} else {
			vars[RunChangeExist] = "true"
			vars[RunChangeMerged] = strconv.FormatBool(ch.Merged)
			vars[RunChangeAbandoned] = strconv.FormatBool(ch.Closed)
			varsFloats[RunChangeDayBefore] = math.Floor(time.Now().Sub(ch.Updated).Hours())
		}
	}

	// If we have a branch in payload, check if it exists on repository branches list
	b, has := vars["git.branch"]
	var exist bool
	if has {
		_, exist = branchesMap[b]
	}
	if has && vcsClient == nil {
		return sdk.NewErrorFrom(sdk.ErrUnknownError, "vcsClient nil but git branch exists")
	}
	vars[RunHasGitBranch] = strconv.FormatBool(has)
	vars[RunGitBranchExist] = strconv.FormatBool(exist)

	vars[RunStatus] = run.Status

	varsFloats[RunDaysBefore] = math.Floor(time.Now().Sub(run.LastModified).Hours() / 24)

	luaCheck.SetVariables(vars)
	luaCheck.SetFloatVariables(varsFloats)
	return nil
}
<|endoftext|>


<commit_before>package dht

import (
	"context"
	"fmt"
	"math"
	"time"

	peer "github.com/libp2p/go-libp2p-peer"
	queue "github.com/libp2p/go-libp2p-peerstore/queue"
)

const (
	// DefaultDialQueueMinParallelism is the default value for the minimum number of worker dial goroutines that will
	// be alive at any time.
	DefaultDialQueueMinParallelism = 6
	// DefaultDialQueueMaxParallelism is the default value for the maximum number of worker dial goroutines that can
	// be alive at any time.
	DefaultDialQueueMaxParallelism = 20
	// DefaultDialQueueMaxIdle is the default value for the period that a worker dial goroutine waits before signalling
	// a worker pool downscaling.
	DefaultDialQueueMaxIdle = 5 * time.Second
	// DefaultDialQueueScalingMutePeriod is the default value for the amount of time to ignore further worker pool
	// scaling events, after one is processed. Its role is to reduce jitter.
	DefaultDialQueueScalingMutePeriod = 1 * time.Second
	// DefaultDialQueueScalingFactor is the default factor by which the current number of workers will be multiplied
	// or divided when upscaling and downscaling events occur, respectively.
	DefaultDialQueueScalingFactor = 1.5
)

type dialQueue struct {
	*dqParams

	nWorkers uint
	out      *queue.ChanQueue

	waitingCh chan waitingCh
	dieCh     chan struct{}
	growCh    chan struct{}
	shrinkCh  chan struct{}
}

type dqParams struct {
	ctx    context.Context
	target string
	dialFn func(context.Context, peer.ID) error
	in     *queue.ChanQueue
	config dqConfig
}

type dqConfig struct {
	// minParallelism is the minimum number of worker dial goroutines that will be alive at any time.
	minParallelism uint
	// maxParallelism is the maximum number of worker dial goroutines that can be alive at any time.
	maxParallelism uint
	// scalingFactor is the factor by which the current number of workers will be multiplied or divided when upscaling
	// and downscaling events occur, respectively.
	scalingFactor float64
	// mutePeriod is the amount of time to ignore further worker pool scaling events, after one is processed.
	// Its role is to reduce jitter.
	mutePeriod time.Duration
	// maxIdle is the period that a worker dial goroutine waits before signalling a worker pool downscaling.
	maxIdle time.Duration
}

// dqDefaultConfig returns the default configuration for dial queues. See const documentation to learn the default values.
func dqDefaultConfig() dqConfig {
	return dqConfig{
		minParallelism: DefaultDialQueueMinParallelism,
		maxParallelism: DefaultDialQueueMaxParallelism,
		scalingFactor:  DefaultDialQueueScalingFactor,
		maxIdle:        DefaultDialQueueMaxIdle,
		mutePeriod:     DefaultDialQueueScalingMutePeriod,
	}
}

func (dqc *dqConfig) validate() error {
	if dqc.minParallelism > dqc.maxParallelism {
		return fmt.Errorf("minParallelism must be below maxParallelism; actual values: min=%d, max=%d",
			dqc.minParallelism, dqc.maxParallelism)
	}
	if dqc.scalingFactor < 1 {
		return fmt.Errorf("scalingFactor must be >= 1; actual value: %f", dqc.scalingFactor)
	}
	return nil
}

type waitingCh struct {
	ch chan<- peer.ID
	ts time.Time
}

// newDialQueue returns an adaptive dial queue that spawns a dynamically sized set of goroutines to preemptively
// stage dials for later handoff to the DHT protocol for RPC. It identifies backpressure on both ends (dial consumers
// and dial producers), and takes compensating action by adjusting the worker pool.
//
// Why? Dialing is expensive. It's orders of magnitude slower than running an RPC on an already-established
// connection, as it requires establishing a TCP connection, multistream handshake, crypto handshake, mux handshake,
// and protocol negotiation.
//
// We start with config.minParallelism number of workers, and scale up and down based on demand and supply of
// dialled peers.
//
// The following events trigger scaling:
// - we scale up when we can't immediately return a successful dial to a new consumer.
// - we scale down when we've been idle for a while waiting for new dial attempts.
// - we scale down when we complete a dial and realise nobody was waiting for it.
//
// Dialler throttling (e.g. FD limit exceeded) is a concern, as we can easily spin up more workers to compensate, and
// end up adding fuel to the fire. Since we have no deterministic way to detect this for now, we hard-limit concurrency
// to config.maxParallelism.
func newDialQueue(params *dqParams) (*dialQueue, error) {
	sq := &dialQueue{
		dqParams:  params,
		nWorkers:  params.config.minParallelism,
		out:       queue.NewChanQueue(params.ctx, queue.NewXORDistancePQ(params.target)),
		growCh:    make(chan struct{}, 1),
		shrinkCh:  make(chan struct{}, 1),
		waitingCh: make(chan waitingCh),
		dieCh:     make(chan struct{}, params.config.maxParallelism),
	}

	for i := 0; i < int(params.config.minParallelism); i++ {
		go sq.worker()
	}
	go sq.control()
	return sq, nil
}

func (dq *dialQueue) control() {
	var (
		dialled        <-chan peer.ID
		waiting        []waitingCh
		lastScalingEvt = time.Now()
	)

	defer func() {
		for _, w := range waiting {
			close(w.ch)
		}
		waiting = nil
	}()

	for {
		// First process any backlog of dial jobs and waiters -- making progress is the priority.
		// This block is copied below; couldn't find a more concise way of doing this.
		select {
		case <-dq.ctx.Done():
			return
		case w := <-dq.waitingCh:
			waiting = append(waiting, w)
			dialled = dq.out.DeqChan
			continue // onto the top.
		case p, ok := <-dialled:
			if !ok {
				return // we're done if the ChanQueue is closed, which happens when the context is closed.
			}
			w := waiting[0]
			log.Debugf("delivering dialled peer to DHT; took %dms.", time.Since(w.ts)/time.Millisecond)
			w.ch <- p
			close(w.ch)
			waiting = waiting[1:]
			if len(waiting) == 0 {
				// no more waiters, so stop consuming dialled jobs.
				dialled = nil
			}
			continue // onto the top.
		default:
			// there's nothing to process, so proceed onto the main select block.
		}

		select {
		case <-dq.ctx.Done():
			return
		case w := <-dq.waitingCh:
			waiting = append(waiting, w)
			dialled = dq.out.DeqChan
		case p, ok := <-dialled:
			if !ok {
				return // we're done if the ChanQueue is closed, which happens when the context is closed.
			}
			w := waiting[0]
			log.Debugf("delivering dialled peer to DHT; took %dms.", time.Since(w.ts)/time.Millisecond)
			w.ch <- p
			close(w.ch)
			waiting = waiting[1:]
			if len(waiting) == 0 {
				// no more waiters, so stop consuming dialled jobs.
				dialled = nil
			}
		case <-dq.growCh:
			if time.Since(lastScalingEvt) < dq.config.mutePeriod {
				continue
			}
			dq.grow()
			lastScalingEvt = time.Now()
		case <-dq.shrinkCh:
			if time.Since(lastScalingEvt) < dq.config.mutePeriod {
				continue
			}
			dq.shrink()
			lastScalingEvt = time.Now()
		}
	}
}

func (dq *dialQueue) Consume() <-chan peer.ID {
	ch := make(chan peer.ID, 1)

	select {
	case p := <-dq.out.DeqChan:
		// short circuit and return a dialled peer if it's immediately available.
		ch <- p
		close(ch)
		return ch
	case <-dq.ctx.Done():
		// return a closed channel with no value if we're done.
		close(ch)
		return ch
	default:
	}

	// we have no finished dials to return, trigger a scale up.
	select {
	case dq.growCh <- struct{}{}:
	default:
	}

	// park the channel until a dialled peer becomes available.
	select {
	case dq.waitingCh <- waitingCh{ch, time.Now()}:
		// all good
	case <-dq.ctx.Done():
		// return a closed channel with no value if we're done.
		close(ch)
	}
	return ch
}

func (dq *dialQueue) grow() {
	// no mutex needed as this is only called from the (single-threaded) control loop.
	defer func(prev uint) {
		if prev == dq.nWorkers {
			return
		}
		log.Debugf("grew dial worker pool: %d => %d", prev, dq.nWorkers)
	}(dq.nWorkers)

	if dq.nWorkers == dq.config.maxParallelism {
		return
	}
	// choosing not to worry about uint wrapping beyond max value.
	target := uint(math.Floor(float64(dq.nWorkers) * dq.config.scalingFactor))
	if target > dq.config.maxParallelism {
		target = dq.config.maxParallelism
	}
	for ; dq.nWorkers < target; dq.nWorkers++ {
		go dq.worker()
	}
}

func (dq *dialQueue) shrink() {
	// no mutex needed as this is only called from the (single-threaded) control loop.
	defer func(prev uint) {
		if prev == dq.nWorkers {
			return
		}
		log.Debugf("shrunk dial worker pool: %d => %d", prev, dq.nWorkers)
	}(dq.nWorkers)

	if dq.nWorkers == dq.config.minParallelism {
		return
	}
	target := uint(math.Floor(float64(dq.nWorkers) / dq.config.scalingFactor))
	if target < dq.config.minParallelism {
		target = dq.config.minParallelism
	}
	// send as many die signals as workers we have to prune.
	for ; dq.nWorkers > target; dq.nWorkers-- {
		select {
		case dq.dieCh <- struct{}{}:
		default:
			log.Debugf("too many die signals queued up.")
		}
	}
}

func (dq *dialQueue) worker() {
	// This idle timer tracks if the environment is slow. If we're waiting to long to acquire a peer to dial,
	// it means that the DHT query is progressing slow and we should shrink the worker pool.
	idleTimer := time.NewTimer(24 * time.Hour) // placeholder init value which will be overridden immediately.
	for {
		// trap exit signals first.
		select {
		case <-dq.ctx.Done():
			return
		case <-dq.dieCh:
			return
		default:
		}

		idleTimer.Stop()
		select {
		case <-idleTimer.C:
		default:
		}
		idleTimer.Reset(dq.config.maxIdle)

		select {
		case <-dq.dieCh:
			return
		case <-dq.ctx.Done():
			return
		case <-idleTimer.C:
			// no new dial requests during our idle period; time to scale down.
		case p := <-dq.in.DeqChan:
			t := time.Now()
			if err := dq.dialFn(dq.ctx, p); err != nil {
				log.Debugf("discarding dialled peer because of error: %v", err)
				continue
			}
			log.Debugf("dialling %v took %dms (as observed by the dht subsystem).", p, time.Since(t)/time.Millisecond)
			waiting := len(dq.waitingCh)
			dq.out.EnqChan <- p
			if waiting > 0 {
				// we have somebody to deliver this value to, so no need to shrink.
				continue
			}
		}

		// scaling down; control only arrives here if the idle timer fires, or if there are no goroutines
		// waiting for the value we just produced.
		select {
		case dq.shrinkCh <- struct{}{}:
		default:
		}
	}
}
<commit_msg>fix a race in dial queue.<commit_after>package dht

import (
	"context"
	"fmt"
	"math"
	"time"

	peer "github.com/libp2p/go-libp2p-peer"
	queue "github.com/libp2p/go-libp2p-peerstore/queue"
)

const (
	// DefaultDialQueueMinParallelism is the default value for the minimum number of worker dial goroutines that will
	// be alive at any time.
	DefaultDialQueueMinParallelism = 6
	// DefaultDialQueueMaxParallelism is the default value for the maximum number of worker dial goroutines that can
	// be alive at any time.
	DefaultDialQueueMaxParallelism = 20
	// DefaultDialQueueMaxIdle is the default value for the period that a worker dial goroutine waits before signalling
	// a worker pool downscaling.
	DefaultDialQueueMaxIdle = 5 * time.Second
	// DefaultDialQueueScalingMutePeriod is the default value for the amount of time to ignore further worker pool
	// scaling events, after one is processed. Its role is to reduce jitter.
	DefaultDialQueueScalingMutePeriod = 1 * time.Second
	// DefaultDialQueueScalingFactor is the default factor by which the current number of workers will be multiplied
	// or divided when upscaling and downscaling events occur, respectively.
	DefaultDialQueueScalingFactor = 1.5
)

type dialQueue struct {
	*dqParams

	nWorkers uint
	out      *queue.ChanQueue

	waitingCh chan waitingCh
	dieCh     chan struct{}
	growCh    chan struct{}
	shrinkCh  chan struct{}
}

type dqParams struct {
	ctx    context.Context
	target string
	dialFn func(context.Context, peer.ID) error
	in     *queue.ChanQueue
	config dqConfig
}

type dqConfig struct {
	// minParallelism is the minimum number of worker dial goroutines that will be alive at any time.
	minParallelism uint
	// maxParallelism is the maximum number of worker dial goroutines that can be alive at any time.
	maxParallelism uint
	// scalingFactor is the factor by which the current number of workers will be multiplied or divided when upscaling
	// and downscaling events occur, respectively.
	scalingFactor float64
	// mutePeriod is the amount of time to ignore further worker pool scaling events, after one is processed.
	// Its role is to reduce jitter.
	mutePeriod time.Duration
	// maxIdle is the period that a worker dial goroutine waits before signalling a worker pool downscaling.
	maxIdle time.Duration
}

// dqDefaultConfig returns the default configuration for dial queues. See const documentation to learn the default values.
func dqDefaultConfig() dqConfig {
	return dqConfig{
		minParallelism: DefaultDialQueueMinParallelism,
		maxParallelism: DefaultDialQueueMaxParallelism,
		scalingFactor:  DefaultDialQueueScalingFactor,
		maxIdle:        DefaultDialQueueMaxIdle,
		mutePeriod:     DefaultDialQueueScalingMutePeriod,
	}
}

func (dqc *dqConfig) validate() error {
	if dqc.minParallelism > dqc.maxParallelism {
		return fmt.Errorf("minParallelism must be below maxParallelism; actual values: min=%d, max=%d",
			dqc.minParallelism, dqc.maxParallelism)
	}
	if dqc.scalingFactor < 1 {
		return fmt.Errorf("scalingFactor must be >= 1; actual value: %f", dqc.scalingFactor)
	}
	return nil
}

type waitingCh struct {
	ch chan<- peer.ID
	ts time.Time
}

// newDialQueue returns an adaptive dial queue that spawns a dynamically sized set of goroutines to preemptively
// stage dials for later handoff to the DHT protocol for RPC. It identifies backpressure on both ends (dial consumers
// and dial producers), and takes compensating action by adjusting the worker pool.
//
// Why? Dialing is expensive. It's orders of magnitude slower than running an RPC on an already-established
// connection, as it requires establishing a TCP connection, multistream handshake, crypto handshake, mux handshake,
// and protocol negotiation.
//
// We start with config.minParallelism number of workers, and scale up and down based on demand and supply of
// dialled peers.
//
// The following events trigger scaling:
// - we scale up when we can't immediately return a successful dial to a new consumer.
// - we scale down when we've been idle for a while waiting for new dial attempts.
// - we scale down when we complete a dial and realise nobody was waiting for it.
//
// Dialler throttling (e.g. FD limit exceeded) is a concern, as we can easily spin up more workers to compensate, and
// end up adding fuel to the fire. Since we have no deterministic way to detect this for now, we hard-limit concurrency
// to config.maxParallelism.
func newDialQueue(params *dqParams) (*dialQueue, error) {
	sq := &dialQueue{
		dqParams:  params,
		nWorkers:  params.config.minParallelism,
		out:       queue.NewChanQueue(params.ctx, queue.NewXORDistancePQ(params.target)),
		growCh:    make(chan struct{}, 1),
		shrinkCh:  make(chan struct{}, 1),
		waitingCh: make(chan waitingCh),
		dieCh:     make(chan struct{}, params.config.maxParallelism),
	}

	for i := 0; i < int(params.config.minParallelism); i++ {
		go sq.worker()
	}
	go sq.control()
	return sq, nil
}

func (dq *dialQueue) control() {
	var (
		dialled        <-chan peer.ID
		waiting        []waitingCh
		lastScalingEvt = time.Now()
	)

	defer func() {
		for _, w := range waiting {
			close(w.ch)
		}
		waiting = nil
	}()

	for {
		// First process any backlog of dial jobs and waiters -- making progress is the priority.
		// This block is copied below; couldn't find a more concise way of doing this.
		select {
		case <-dq.ctx.Done():
			return
		case w := <-dq.waitingCh:
			waiting = append(waiting, w)
			dialled = dq.out.DeqChan
			continue // onto the top.
		case p, ok := <-dialled:
			if !ok {
				return // we're done if the ChanQueue is closed, which happens when the context is closed.
			}
			w := waiting[0]
			log.Debugf("delivering dialled peer to DHT; took %dms.", time.Since(w.ts)/time.Millisecond)
			w.ch <- p
			close(w.ch)
			waiting = waiting[1:]
			if len(waiting) == 0 {
				// no more waiters, so stop consuming dialled jobs.
				dialled = nil
			}
			continue // onto the top.
		default:
			// there's nothing to process, so proceed onto the main select block.
		}

		select {
		case <-dq.ctx.Done():
			return
		case w := <-dq.waitingCh:
			waiting = append(waiting, w)
			dialled = dq.out.DeqChan
		case p, ok := <-dialled:
			if !ok {
				return // we're done if the ChanQueue is closed, which happens when the context is closed.
			}
			w := waiting[0]
			log.Debugf("delivering dialled peer to DHT; took %dms.", time.Since(w.ts)/time.Millisecond)
			w.ch <- p
			close(w.ch)
			waiting = waiting[1:]
			if len(waiting) == 0 {
				// no more waiters, so stop consuming dialled jobs.
				dialled = nil
			}
		case <-dq.growCh:
			if time.Since(lastScalingEvt) < dq.config.mutePeriod {
				continue
			}
			dq.grow()
			lastScalingEvt = time.Now()
		case <-dq.shrinkCh:
			if time.Since(lastScalingEvt) < dq.config.mutePeriod {
				continue
			}
			dq.shrink()
			lastScalingEvt = time.Now()
		}
	}
}

func (dq *dialQueue) Consume() <-chan peer.ID {
	ch := make(chan peer.ID, 1)

	select {
	case p, ok := <-dq.out.DeqChan:
		// short circuit and return a dialled peer if it's immediately available, or abort if DeqChan is closed.
		if ok {
			ch <- p
		}
		close(ch)
		return ch
	case <-dq.ctx.Done():
		// return a closed channel with no value if we're done.
		close(ch)
		return ch
	default:
	}

	// we have no finished dials to return, trigger a scale up.
	select {
	case dq.growCh <- struct{}{}:
	default:
	}

	// park the channel until a dialled peer becomes available.
	select {
	case dq.waitingCh <- waitingCh{ch, time.Now()}:
		// all good
	case <-dq.ctx.Done():
		// return a closed channel with no value if we're done.
		close(ch)
	}
	return ch
}

func (dq *dialQueue) grow() {
	// no mutex needed as this is only called from the (single-threaded) control loop.
	defer func(prev uint) {
		if prev == dq.nWorkers {
			return
		}
		log.Debugf("grew dial worker pool: %d => %d", prev, dq.nWorkers)
	}(dq.nWorkers)

	if dq.nWorkers == dq.config.maxParallelism {
		return
	}
	// choosing not to worry about uint wrapping beyond max value.
	target := uint(math.Floor(float64(dq.nWorkers) * dq.config.scalingFactor))
	if target > dq.config.maxParallelism {
		target = dq.config.maxParallelism
	}
	for ; dq.nWorkers < target; dq.nWorkers++ {
		go dq.worker()
	}
}

func (dq *dialQueue) shrink() {
	// no mutex needed as this is only called from the (single-threaded) control loop.
	defer func(prev uint) {
		if prev == dq.nWorkers {
			return
		}
		log.Debugf("shrunk dial worker pool: %d => %d", prev, dq.nWorkers)
	}(dq.nWorkers)

	if dq.nWorkers == dq.config.minParallelism {
		return
	}
	target := uint(math.Floor(float64(dq.nWorkers) / dq.config.scalingFactor))
	if target < dq.config.minParallelism {
		target = dq.config.minParallelism
	}
	// send as many die signals as workers we have to prune.
	for ; dq.nWorkers > target; dq.nWorkers-- {
		select {
		case dq.dieCh <- struct{}{}:
		default:
			log.Debugf("too many die signals queued up.")
		}
	}
}

func (dq *dialQueue) worker() {
	// This idle timer tracks if the environment is slow. If we're waiting to long to acquire a peer to dial,
	// it means that the DHT query is progressing slow and we should shrink the worker pool.
	idleTimer := time.NewTimer(24 * time.Hour) // placeholder init value which will be overridden immediately.
	for {
		// trap exit signals first.
		select {
		case <-dq.ctx.Done():
			return
		case <-dq.dieCh:
			return
		default:
		}

		idleTimer.Stop()
		select {
		case <-idleTimer.C:
		default:
		}
		idleTimer.Reset(dq.config.maxIdle)

		select {
		case <-dq.dieCh:
			return
		case <-dq.ctx.Done():
			return
		case <-idleTimer.C:
			// no new dial requests during our idle period; time to scale down.
		case p := <-dq.in.DeqChan:
			t := time.Now()
			if err := dq.dialFn(dq.ctx, p); err != nil {
				log.Debugf("discarding dialled peer because of error: %v", err)
				continue
			}
			log.Debugf("dialling %v took %dms (as observed by the dht subsystem).", p, time.Since(t)/time.Millisecond)
			waiting := len(dq.waitingCh)
			dq.out.EnqChan <- p
			if waiting > 0 {
				// we have somebody to deliver this value to, so no need to shrink.
				continue
			}
		}

		// scaling down; control only arrives here if the idle timer fires, or if there are no goroutines
		// waiting for the value we just produced.
		select {
		case dq.shrinkCh <- struct{}{}:
		default:
		}
	}
}
<|endoftext|>


<commit_before>package handler

import (
	"bufio"
	"fmt"
	"github.com/realglobe-Inc/go-lib-rg/erro"
	"github.com/realglobe-Inc/go-lib-rg/rglog/level"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
)

type FormatEntry struct {
	Date time.Time
	File string
	Line int
	Lv   level.Level
	Args []interface{}
}

// てきとうなポインタを Flush のトリガにする。
// ポインタ自体で判断するので、中身を保護する必要は無い。
var flushTrigger = &FormatEntry{}

func FlushTrigger() *FormatEntry {
	return flushTrigger
}

type goHandler struct {
	sync.Mutex
	lv level.Level

	done  chan bool
	queue chan *FormatEntry
}

func (hndl *goHandler) Output(depth int, lv level.Level, v ...interface{}) {
	hndl.Lock()
	if lv > hndl.lv {
		hndl.Unlock()
		return
	}
	hndl.Unlock()

	now := time.Now()
	var file string
	var line int
	var ok bool
	_, file, line, ok = runtime.Caller(depth + 1)
	if !ok {
		file = "???"
		line = 0
	}

	hndl.queue <- &FormatEntry{now, file, line, lv, v}
}

func (hndl *goHandler) SetLevel(lv level.Level) {
	hndl.Lock()
	defer hndl.Unlock()

	hndl.lv = lv
}

func (hndl *goHandler) Flush() {
	hndl.queue <- FlushTrigger()
	<-hndl.done
}

func NewGoHandler(goFunc func(done chan bool, queue chan *FormatEntry), queueCapacity int) Handler {
	done := make(chan bool)
	queue := make(chan *FormatEntry, queueCapacity)
	go goFunc(done, queue)
	return &goHandler{done: done, queue: queue}
}

const defaultQueueCapacity = 8192

func NewRotateHandler(path string, limit int64, num int) (Handler, error) {
	return NewRotateHandlerUsing(path, limit, num, defaultQueueCapacity, &simpleFormatter{})
}

const errThreshold = 5                    // 再試行する限度。
const coolDownDuration = time.Millisecond // 異常発生時に空ける間隔。

func NewRotateHandlerUsing(path string, limit int64, num, queueCapacity int, formatter Formatter) (Handler, error) {
	// ファイルチェック。
	file, err := os.OpenFile(path, os.O_RDWR|os.O_APPEND|os.O_CREATE, logPerm)
	if err != nil {
		return nil, erro.Wrap(err)
	}
	if err := file.Close(); err != nil {
		return nil, erro.Wrap(err)
	}

	return NewGoHandler(func(done chan bool, queue chan *FormatEntry) {
		for errCount := 0; errCount < errThreshold; {

			// 異常が発生してたら、ちょっと落ち着く。
			if errCount > 0 {
				time.Sleep(coolDownDuration)
			}

			// ファイルを開く。
			file, err := os.OpenFile(path, os.O_RDWR|os.O_APPEND|os.O_CREATE, logPerm)
			if err != nil {
				fmt.Fprintln(os.Stderr, erro.Wrap(err))
				errCount++
				continue
			}

			fi, err := file.Stat()
			if err != nil {
				fmt.Fprintln(os.Stderr, erro.Wrap(err))
				errCount++
				file.Close()
				continue
			}

			size := fi.Size()
			writer := bufio.NewWriter(file)

			// ログを取る。
			for size <= limit { // 最大 1 エントリ分はみ出す。でも、limit == 0 でも動く。
				ent := <-queue

				if ent == FlushTrigger() {
					err := writer.Flush()
					if err != nil {
						fmt.Fprintln(os.Stderr, erro.Wrap(err))
						done <- (err == nil)
						break
					}
					done <- (err == nil)
				}

				buff := formatter.Format(ent.Date, ent.File, ent.Line, ent.Lv, ent.Args...)
				_, err := writer.Write(buff)
				if err != nil {
					fmt.Fprintln(os.Stderr, erro.Wrap(err))
					break
				}

				size += int64(len(buff))
			}

			// ファイルを閉じる。
			if err := writer.Flush(); err != nil {
				fmt.Fprintln(os.Stderr, erro.Wrap(err))
				errCount++
				file.Close()
				continue
			}

			if err := file.Close(); err != nil {
				fmt.Fprintln(os.Stderr, erro.Wrap(err))
				errCount++
				continue
			}

			if size < limit { // ログを取ってるときに異常が発生してた。
				errCount++
				continue
			}

			// ファイルを回す。
			if err := rotateFile(path, num); err != nil {
				fmt.Fprintln(os.Stderr, erro.Wrap(err))
				errCount++
				continue
			}

			errCount = 0
		}

		// 異常でループを抜けたら、デッドロック防止処理だけする。
		for {
			ent := <-queue
			if ent == FlushTrigger() {
				done <- false
			}
		}
	}, queueCapacity), nil
}

func rotateFile(path string, num int) error {
	var n int

	for n = 1; n < num; n++ {
		_, err := os.Stat(path + "." + strconv.Itoa(n))
		if err == nil {
			continue
		} else if os.IsNotExist(err) {
			break
		} else {
			return erro.Wrap(err)
		}
	}
	n--

	// .{n} が残す中で一番最後。

	// .{i} を .{i+1} に。
	for ; n > 0; n-- {
		from := path + "." + strconv.Itoa(n)
		to := path + "." + strconv.Itoa(n+1)
		if err := os.Rename(from, to); err != nil {
			return erro.Wrap(err)
		}
	}

	// 最新版 を .1 に。
	if err := os.Rename(path, path+".1"); err != nil {
		return erro.Wrap(err)
	}

	return nil
}
<commit_msg>パニックが起きても周りを巻き込まないようにした<commit_after>package handler

import (
	"bufio"
	"fmt"
	"github.com/realglobe-Inc/go-lib-rg/erro"
	"github.com/realglobe-Inc/go-lib-rg/rglog/level"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
)

type FormatEntry struct {
	Date time.Time
	File string
	Line int
	Lv   level.Level
	Args []interface{}
}

// てきとうなポインタを Flush のトリガにする。
// ポインタ自体で判断するので、中身を保護する必要は無い。
var flushTrigger = &FormatEntry{}

func FlushTrigger() *FormatEntry {
	return flushTrigger
}

type goHandler struct {
	sync.Mutex
	lv level.Level

	done  chan bool
	queue chan *FormatEntry
}

func (hndl *goHandler) Output(depth int, lv level.Level, v ...interface{}) {
	hndl.Lock()
	if lv > hndl.lv {
		hndl.Unlock()
		return
	}
	hndl.Unlock()

	now := time.Now()
	var file string
	var line int
	var ok bool
	_, file, line, ok = runtime.Caller(depth + 1)
	if !ok {
		file = "???"
		line = 0
	}

	hndl.queue <- &FormatEntry{now, file, line, lv, v}
}

func (hndl *goHandler) SetLevel(lv level.Level) {
	hndl.Lock()
	defer hndl.Unlock()

	hndl.lv = lv
}

func (hndl *goHandler) Flush() {
	hndl.queue <- FlushTrigger()
	<-hndl.done
}

func NewGoHandler(goFunc func(done chan bool, queue chan *FormatEntry), queueCapacity int) Handler {
	done := make(chan bool)
	queue := make(chan *FormatEntry, queueCapacity)
	go goFunc(done, queue)
	return &goHandler{done: done, queue: queue}
}

const defaultQueueCapacity = 8192

func NewRotateHandler(path string, limit int64, num int) (Handler, error) {
	return NewRotateHandlerUsing(path, limit, num, defaultQueueCapacity, &simpleFormatter{})
}

const errThreshold = 5                    // 再試行する限度。
const coolDownDuration = time.Millisecond // 異常発生時に空ける間隔。

func NewRotateHandlerUsing(path string, limit int64, num, queueCapacity int, formatter Formatter) (Handler, error) {
	// ロックファイルをつくったほうが良いが、OS 依存なので今は止めとく。

	// ファイルチェック。
	file, err := os.OpenFile(path, os.O_RDWR|os.O_APPEND|os.O_CREATE, logPerm)
	if err != nil {
		return nil, erro.Wrap(err)
	}
	if err := file.Close(); err != nil {
		return nil, erro.Wrap(err)
	}

	return NewGoHandler(func(done chan bool, queue chan *FormatEntry) {
		func() {
			defer func() {
				if err := recover(); err != nil {
					fmt.Fprintln(os.Stderr, err)
				}
			}()

			for errCount := 0; errCount < errThreshold; {

				// 異常が発生してたら、ちょっと落ち着く。
				if errCount > 0 {
					time.Sleep(coolDownDuration)
				}

				// ファイルを開く。
				file, err := os.OpenFile(path, os.O_RDWR|os.O_APPEND|os.O_CREATE, logPerm)
				if err != nil {
					fmt.Fprintln(os.Stderr, erro.Wrap(err))
					errCount++
					continue
				}

				fi, err := file.Stat()
				if err != nil {
					fmt.Fprintln(os.Stderr, erro.Wrap(err))
					errCount++
					file.Close()
					continue
				}

				size := fi.Size()
				writer := bufio.NewWriter(file)

				// ログを取る。
				for size <= limit { // 最大 1 エントリ分はみ出す。でも、limit == 0 でも動く。
					ent := <-queue

					if ent == FlushTrigger() {
						err := writer.Flush()
						if err != nil {
							fmt.Fprintln(os.Stderr, erro.Wrap(err))
							done <- (err == nil)
							break
						}
						done <- (err == nil)
					}

					buff := formatter.Format(ent.Date, ent.File, ent.Line, ent.Lv, ent.Args...)
					_, err := writer.Write(buff)
					if err != nil {
						fmt.Fprintln(os.Stderr, erro.Wrap(err))
						break
					}

					size += int64(len(buff))
				}

				// ファイルを閉じる。
				if err := writer.Flush(); err != nil {
					fmt.Fprintln(os.Stderr, erro.Wrap(err))
					errCount++
					file.Close()
					continue
				}

				if err := file.Close(); err != nil {
					fmt.Fprintln(os.Stderr, erro.Wrap(err))
					errCount++
					continue
				}

				if size < limit { // ログを取ってるときに異常が発生してた。
					errCount++
					continue
				}

				// ファイルを回す。
				if err := rotateFile(path, num); err != nil {
					fmt.Fprintln(os.Stderr, erro.Wrap(err))
					errCount++
					continue
				}

				errCount = 0
			}

			fmt.Fprintln(os.Stderr, "Too many errors.")
		}()

		fmt.Fprintln(os.Stderr, "Logging into ", path, " was aborted.")

		// 異常でループを抜けたら、デッドロック防止処理だけする。
		for {
			ent := <-queue
			if ent == FlushTrigger() {
				done <- false
			}
		}
	}, queueCapacity), nil
}

func rotateFile(path string, num int) error {
	var n int

	for n = 1; n < num; n++ {
		_, err := os.Stat(path + "." + strconv.Itoa(n))
		if err == nil {
			continue
		} else if os.IsNotExist(err) {
			break
		} else {
			return erro.Wrap(err)
		}
	}
	n--

	// .{n} が残す中で一番最後。

	// .{i} を .{i+1} に。
	for ; n > 0; n-- {
		from := path + "." + strconv.Itoa(n)
		to := path + "." + strconv.Itoa(n+1)
		if err := os.Rename(from, to); err != nil {
			return erro.Wrap(err)
		}
	}

	// 最新版 を .1 に。
	if err := os.Rename(path, path+".1"); err != nil {
		return erro.Wrap(err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Package klient provides an instance and abstraction to a remote klient kite.
// It is used to easily call methods of a klient kite
package klient

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"koding/kite-handler/command"
	"koding/kites/klient/usage"

	"github.com/koding/kite"
	"github.com/koding/kite/protocol"
	"github.com/koding/logging"
)

var ErrDialingFailed = errors.New("Dialing klient failed.")

// KlientPool represents a pool of connected klients
type KlientPool struct {
	kite    *kite.Kite
	klients map[string]*Klient
	log     logging.Logger
	sync.Mutex
}

// Klient represents a remote klient instance
type Klient struct {
	client   *kite.Client
	kite     *kite.Kite
	Username string
}

func NewPool(k *kite.Kite) *KlientPool {
	return &KlientPool{
		kite:    k,
		klients: make(map[string]*Klient),
		log:     logging.NewLogger("klientpool"),
	}
}

// Get returns a ready to use and connected klient from the pool.
func (k *KlientPool) Get(queryString string) (*Klient, error) {
	var klient *Klient
	var ok bool
	var err error

	k.Lock()
	defer k.Unlock()

	klient, ok = k.klients[queryString]
	if !ok {
		klient, err = Connect(k.kite, queryString)
		if err != nil {
			return nil, err
		}

		k.log.Info("creating new klient connection to %s", queryString)
		k.klients[queryString] = klient

		// remove from the pool if we loose the connection
		klient.client.OnDisconnect(func() {
			k.log.Info("klient %s disconnected. removing from the pool", queryString)
			k.Delete(queryString)
			klient.Close()
		})
	} else {
		k.log.Debug("fetching already connected klient (%s) from pool", queryString)
	}

	return klient, nil
}

// Delete removes the klient with the given queryString from the pool
func (k *KlientPool) Delete(queryString string) {
	k.Lock()
	defer k.Unlock()

	delete(k.klients, queryString)
}

// Exists checks whether the given queryString exists in Kontrol or not
func Exists(k *kite.Kite, queryString string) error {
	query, err := protocol.KiteFromString(queryString)
	if err != nil {
		return err
	}

	k.Log.Debug("Querying for Klient: %s", queryString)

	// an error indicates a non existing klient or another error.
	_, err = k.GetKites(query.Query())
	if err != nil {
		return err
	}

	return nil
}

// Connect returns a new connected klient instance to the given queryString. The
// klient is ready to use. It's connected and needs to be closed once the task
// is finished with it.
func Connect(k *kite.Kite, queryString string) (*Klient, error) {
	query, err := protocol.KiteFromString(queryString)
	if err != nil {
		return nil, err
	}

	k.Log.Debug("Querying for Klient: %s", queryString)

	kites, err := k.GetKites(query.Query())
	if err != nil {
		return nil, err
	}

	remoteKite := kites[0]
	if err := remoteKite.Dial(); err != nil {
		return nil, ErrDialingFailed
	}

	// klient connection is ready now
	return &Klient{
		kite:     k,
		client:   remoteKite,
		Username: remoteKite.Username,
	}, nil

}

// Klient returns a new connected klient instance to the given queryString. The
// klient is ready to use. It's tries to connect for the given timeout duration
func NewWithTimeout(k *kite.Kite, queryString string, t time.Duration) (*Klient, error) {
	timeout := time.After(t)

	k.Log.Debug("Querying for Klient: %s", queryString)
	for {
		select {
		case <-time.Tick(time.Second * 2):
			if klient, err := Connect(k, queryString); err == nil {
				return klient, nil
			}

		case <-timeout:
			return nil, ErrDialingFailed
		}
	}
}

func (k *Klient) Close() {
	k.client.Close()
}

// Usage calls the usage method of remote and get's the result back
func (k *Klient) Usage() (*usage.Usage, error) {
	resp, err := k.client.Tell(usage.MethodName)
	if err != nil {
		return nil, err
	}

	var usg *usage.Usage
	if err := resp.Unmarshal(&usg); err != nil {
		return nil, err
	}

	return usg, nil
}

// Exec runs a shell command on remote klient machinek
func (k *Klient) Exec(cmd string) (*command.Output, error) {
	var params = struct {
		Command string
	}{
		Command: cmd,
	}

	resp, err := k.client.Tell("exec", params)
	if err != nil {
		return nil, err
	}

	var out *command.Output
	if err := resp.Unmarshal(&out); err != nil {
		return nil, err
	}

	return out, nil
}

// Ping checks if the given klient response with "pong" to the "ping" we send.
// A nil error means a successfull pong result.
func (k *Klient) Ping() error {
	resp, err := k.client.TellWithTimeout("kite.ping", 10*time.Second)
	if err != nil {
		return err
	}

	out, err := resp.String()
	if err != nil {
		return err
	}

	if out == "pong" {
		return nil
	}

	return fmt.Errorf("wrong response %s", out)
}
<commit_msg>kloud: klient connections are short living, make them low<commit_after>// Package klient provides an instance and abstraction to a remote klient kite.
// It is used to easily call methods of a klient kite
package klient

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"koding/kite-handler/command"
	"koding/kites/klient/usage"

	"github.com/koding/kite"
	"github.com/koding/kite/protocol"
	"github.com/koding/logging"
)

var ErrDialingFailed = errors.New("Dialing klient failed.")

// KlientPool represents a pool of connected klients
type KlientPool struct {
	kite    *kite.Kite
	klients map[string]*Klient
	log     logging.Logger
	sync.Mutex
}

// Klient represents a remote klient instance
type Klient struct {
	client   *kite.Client
	kite     *kite.Kite
	Username string
}

func NewPool(k *kite.Kite) *KlientPool {
	return &KlientPool{
		kite:    k,
		klients: make(map[string]*Klient),
		log:     logging.NewLogger("klientpool"),
	}
}

// Get returns a ready to use and connected klient from the pool.
func (k *KlientPool) Get(queryString string) (*Klient, error) {
	var klient *Klient
	var ok bool
	var err error

	k.Lock()
	defer k.Unlock()

	klient, ok = k.klients[queryString]
	if !ok {
		klient, err = Connect(k.kite, queryString)
		if err != nil {
			return nil, err
		}

		k.log.Info("creating new klient connection to %s", queryString)
		k.klients[queryString] = klient

		// remove from the pool if we loose the connection
		klient.client.OnDisconnect(func() {
			k.log.Info("klient %s disconnected. removing from the pool", queryString)
			k.Delete(queryString)
			klient.Close()
		})
	} else {
		k.log.Debug("fetching already connected klient (%s) from pool", queryString)
	}

	return klient, nil
}

// Delete removes the klient with the given queryString from the pool
func (k *KlientPool) Delete(queryString string) {
	k.Lock()
	defer k.Unlock()

	delete(k.klients, queryString)
}

// Exists checks whether the given queryString exists in Kontrol or not
func Exists(k *kite.Kite, queryString string) error {
	query, err := protocol.KiteFromString(queryString)
	if err != nil {
		return err
	}

	k.Log.Debug("Querying for Klient: %s", queryString)

	// an error indicates a non existing klient or another error.
	_, err = k.GetKites(query.Query())
	if err != nil {
		return err
	}

	return nil
}

// Connect returns a new connected klient instance to the given queryString. The
// klient is ready to use. It's connected and needs to be closed once the task
// is finished with it.
func Connect(k *kite.Kite, queryString string) (*Klient, error) {
	query, err := protocol.KiteFromString(queryString)
	if err != nil {
		return nil, err
	}

	k.Log.Debug("Querying for Klient: %s", queryString)

	kites, err := k.GetKites(query.Query())
	if err != nil {
		return nil, err
	}

	remoteKite := kites[0]
	remoteKite.ReadBufferSize = 512
	remoteKite.WriteBufferSize = 512
	if err := remoteKite.Dial(); err != nil {
		return nil, ErrDialingFailed
	}

	// klient connection is ready now
	return &Klient{
		kite:     k,
		client:   remoteKite,
		Username: remoteKite.Username,
	}, nil

}

// Klient returns a new connected klient instance to the given queryString. The
// klient is ready to use. It's tries to connect for the given timeout duration
func NewWithTimeout(k *kite.Kite, queryString string, t time.Duration) (*Klient, error) {
	timeout := time.After(t)

	k.Log.Debug("Querying for Klient: %s", queryString)
	for {
		select {
		case <-time.Tick(time.Second * 2):
			if klient, err := Connect(k, queryString); err == nil {
				return klient, nil
			}

		case <-timeout:
			return nil, ErrDialingFailed
		}
	}
}

func (k *Klient) Close() {
	k.client.Close()
}

// Usage calls the usage method of remote and get's the result back
func (k *Klient) Usage() (*usage.Usage, error) {
	resp, err := k.client.Tell(usage.MethodName)
	if err != nil {
		return nil, err
	}

	var usg *usage.Usage
	if err := resp.Unmarshal(&usg); err != nil {
		return nil, err
	}

	return usg, nil
}

// Exec runs a shell command on remote klient machinek
func (k *Klient) Exec(cmd string) (*command.Output, error) {
	var params = struct {
		Command string
	}{
		Command: cmd,
	}

	resp, err := k.client.Tell("exec", params)
	if err != nil {
		return nil, err
	}

	var out *command.Output
	if err := resp.Unmarshal(&out); err != nil {
		return nil, err
	}

	return out, nil
}

// Ping checks if the given klient response with "pong" to the "ping" we send.
// A nil error means a successfull pong result.
func (k *Klient) Ping() error {
	resp, err := k.client.TellWithTimeout("kite.ping", 10*time.Second)
	if err != nil {
		return err
	}

	out, err := resp.String()
	if err != nil {
		return err
	}

	if out == "pong" {
		return nil
	}

	return fmt.Errorf("wrong response %s", out)
}
<|endoftext|>


<commit_before>package libreofficekit

// #cgo CFLAGS: -I /usr/include/LibreOfficeKit/ -D LOK_USE_UNSTABLE_API
// #cgo LDFLAGS: -ldl
// #include <stdlib.h>
// #include "LibreOfficeKit/LibreOfficeKitInit.h"
// #include "LibreOfficeKit/LibreOfficeKit.h"
/*
typedef void (*voidFunc) ();
typedef int (*intFunc) ();
typedef char* (*charFunc) ();
void destroy_bridge(voidFunc f, void* handle) {
	return f(handle);
};
LibreOfficeKitDocument* document_load_bridge(voidFunc f,
											 LibreOfficeKit* pThis,
											 const char* pURL) {
	f(pThis, pURL);
};
char* get_error_bridge(charFunc f, LibreOfficeKit* pThis) {
	return f(pThis);
};
int get_document_type_bridge(intFunc f, LibreOfficeKitDocument* pThis) {
	return f(pThis);
};
int document_save_bridge(intFunc f,
						LibreOfficeKitDocument* pThis,
						const char* pUrl,
						const char* pFormat,
						const char* pFilterOptions) {
	return f(pThis, pUrl, pFormat, pFilterOptions);
};
*/
import "C"
import (
	"fmt"
	"unsafe"
)

type Office struct {
	handle *C.struct__LibreOfficeKit
}

func NewOffice(path string) (*Office, error) {
	office := new(Office)

	c_path := C.CString(path)
	defer C.free(unsafe.Pointer(c_path))

	lokit := C.lok_init(c_path)
	if lokit == nil {
		return nil, fmt.Errorf("Failed to initialize LibreOfficeKit with path: '%s'", path)
	}

	office.handle = lokit

	return office, nil

}

func (self *Office) Close() {
	selfPointer := unsafe.Pointer(self.handle)
	C.destroy_bridge(self.handle.pClass.destroy, selfPointer)
}

func (self *Office) GetError() string {
	return C.GoString(C.get_error_bridge(self.handle.pClass.getError, self.handle))
}

func (self *Office) LoadDocument(path string) (*Document, error) {
	document := new(Document)
	c_path := C.CString(path)
	defer C.free(unsafe.Pointer(c_path))
	handle := C.document_load_bridge(self.handle.pClass.documentLoad, self.handle, c_path)
	if handle == nil {
		return nil, fmt.Errorf("Failed to load document")
	}
	document.handle = handle
	return document, nil
}

const (
	TextDocument = iota
	SpreadsheetDocument
	PresentationDocument
	DrawingDocument
	OtherDocument
)

type Document struct {
	handle *C.struct__LibreOfficeKitDocument
}

func (self *Document) Close() {
	selfPointer := unsafe.Pointer(self.handle)
	C.destroy_bridge(self.handle.pClass.destroy, selfPointer)
}

func (self *Document) GetType() int {
	return int(C.get_document_type_bridge(self.handle.pClass.getDocumentType, self.handle))
}

func (self *Document) SaveAs(path string, format string, filter string) error {
	c_path := C.CString(path)
	defer C.free(unsafe.Pointer(c_path))
	c_format := C.CString(format)
	defer C.free(unsafe.Pointer(c_format))
	c_filter := C.CString(filter)
	defer C.free(unsafe.Pointer(c_filter))
	status := C.document_save_bridge(self.handle.pClass.saveAs, self.handle, c_path, c_format, c_filter)
	if status != 0 {
		return fmt.Errorf("Failed to save document")
	} else {
		return nil
	}
}
<commit_msg>Add getDocumentSize support<commit_after>package libreofficekit

// #cgo CFLAGS: -I /usr/include/LibreOfficeKit/ -D LOK_USE_UNSTABLE_API
// #cgo LDFLAGS: -ldl
// #include <stdlib.h>
// #include "LibreOfficeKit/LibreOfficeKitInit.h"
// #include "LibreOfficeKit/LibreOfficeKit.h"
/*
typedef void (*voidFunc) ();
typedef int (*intFunc) ();
typedef char* (*charFunc) ();
void destroy_bridge(voidFunc f, void* handle) {
	return f(handle);
};
LibreOfficeKitDocument* document_load_bridge(voidFunc f,
											 LibreOfficeKit* pThis,
											 const char* pURL) {
	f(pThis, pURL);
};
char* get_error_bridge(charFunc f, LibreOfficeKit* pThis) {
	return f(pThis);
};
int get_document_type_bridge(intFunc f, LibreOfficeKitDocument* pThis) {
	return f(pThis);
};
void get_document_size_bridge(voidFunc f,
							LibreOfficeKitDocument* pThis,
							long* pWidth,
							long* pHeight) {
	return f(pThis, pWidth, pHeight);
};
int document_save_bridge(intFunc f,
						LibreOfficeKitDocument* pThis,
						const char* pUrl,
						const char* pFormat,
						const char* pFilterOptions) {
	return f(pThis, pUrl, pFormat, pFilterOptions);
};
*/
import "C"
import (
	"fmt"
	"unsafe"
)

type Office struct {
	handle *C.struct__LibreOfficeKit
}

func NewOffice(path string) (*Office, error) {
	office := new(Office)

	c_path := C.CString(path)
	defer C.free(unsafe.Pointer(c_path))

	lokit := C.lok_init(c_path)
	if lokit == nil {
		return nil, fmt.Errorf("Failed to initialize LibreOfficeKit with path: '%s'", path)
	}

	office.handle = lokit

	return office, nil

}

func (self *Office) Close() {
	selfPointer := unsafe.Pointer(self.handle)
	C.destroy_bridge(self.handle.pClass.destroy, selfPointer)
}

func (self *Office) GetError() string {
	return C.GoString(C.get_error_bridge(self.handle.pClass.getError, self.handle))
}

func (self *Document) GetSize() (int, int) {
	width := C.long(0)
	heigth := C.long(0)
	C.get_document_size_bridge(self.handle.pClass.getDocumentSize, self.handle, &width, &heigth)
	return int(width), int(heigth)
}

func (self *Office) LoadDocument(path string) (*Document, error) {
	document := new(Document)
	c_path := C.CString(path)
	defer C.free(unsafe.Pointer(c_path))
	handle := C.document_load_bridge(self.handle.pClass.documentLoad, self.handle, c_path)
	if handle == nil {
		return nil, fmt.Errorf("Failed to load document")
	}
	document.handle = handle
	return document, nil
}

const (
	TextDocument = iota
	SpreadsheetDocument
	PresentationDocument
	DrawingDocument
	OtherDocument
)

type Document struct {
	handle *C.struct__LibreOfficeKitDocument
}

func (self *Document) Close() {
	selfPointer := unsafe.Pointer(self.handle)
	C.destroy_bridge(self.handle.pClass.destroy, selfPointer)
}

func (self *Document) GetType() int {
	return int(C.get_document_type_bridge(self.handle.pClass.getDocumentType, self.handle))
}

func (self *Document) SaveAs(path string, format string, filter string) error {
	c_path := C.CString(path)
	defer C.free(unsafe.Pointer(c_path))
	c_format := C.CString(format)
	defer C.free(unsafe.Pointer(c_format))
	c_filter := C.CString(filter)
	defer C.free(unsafe.Pointer(c_filter))
	status := C.document_save_bridge(self.handle.pClass.saveAs, self.handle, c_path, c_format, c_filter)
	if status != 0 {
		return fmt.Errorf("Failed to save document")
	} else {
		return nil
	}
}
<|endoftext|>


<commit_before>package nat

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	manet "QmU5s159q8cZuM1f9Vqti4LHu6y8zyVc5dxv2py81sdp6Q/go-multiaddr-net"
	ma "QmbWxL1aXQhBjc1XGjGF1f2KGBMCBYSuT2ThA8YXnXJK83/go-multiaddr"

	goprocess "QmSir6qPL1tjuxd8LkR8VZq6v625ExAUVs2eCLeqQuaPGU/goprocess"
	periodic "QmSir6qPL1tjuxd8LkR8VZq6v625ExAUVs2eCLeqQuaPGU/goprocess/periodic"
	notifier "QmUtEiB6DmXs7eLJiwS9YFyTAtptqzaWutxCsjHy7UKEgo/go-notifier"
	logging "QmWRypnfEwrgH4k93KEHN5hng7VjKYkWmzDYRuTZeh2Mgh/go-log"
	nat "github.com/fd/go-nat"
)

var (
	// ErrNoMapping signals no mapping exists for an address
	ErrNoMapping = errors.New("mapping not established")
)

var log = logging.Logger("nat")

// MappingDuration is a default port mapping duration.
// Port mappings are renewed every (MappingDuration / 3)
const MappingDuration = time.Second * 60

// CacheTime is the time a mapping will cache an external address for
const CacheTime = time.Second * 15

// DiscoverNAT looks for a NAT device in the network and
// returns an object that can manage port mappings.
func DiscoverNAT() *NAT {
	nat, err := nat.DiscoverGateway()
	if err != nil {
		log.Debug("DiscoverGateway error:", err)
		return nil
	}
	addr, err := nat.GetDeviceAddress()
	if err != nil {
		log.Debug("DiscoverGateway address error:", err)
	} else {
		log.Debug("DiscoverGateway address:", addr)
	}
	return newNAT(nat)
}

// NAT is an object that manages address port mappings in
// NATs (Network Address Translators). It is a long-running
// service that will periodically renew port mappings,
// and keep an up-to-date list of all the external addresses.
type NAT struct {
	nat  nat.NAT
	proc goprocess.Process // manages nat mappings lifecycle

	mappingmu sync.RWMutex // guards mappings
	mappings  map[*mapping]struct{}

	Notifier
}

func newNAT(realNAT nat.NAT) *NAT {
	return &NAT{
		nat:      realNAT,
		proc:     goprocess.WithParent(goprocess.Background()),
		mappings: make(map[*mapping]struct{}),
	}
}

// Close shuts down all port mappings. NAT can no longer be used.
func (nat *NAT) Close() error {
	return nat.proc.Close()
}

// Process returns the nat's life-cycle manager, for making it listen
// to close signals.
func (nat *NAT) Process() goprocess.Process {
	return nat.proc
}

// Notifier is an object that assists NAT in notifying listeners.
// It is implemented using thirdparty/notifier
type Notifier struct {
	n notifier.Notifier
}

func (n *Notifier) notifyAll(notify func(n Notifiee)) {
	n.n.NotifyAll(func(n notifier.Notifiee) {
		notify(n.(Notifiee))
	})
}

// Notify signs up notifiee to listen to NAT events.
func (n *Notifier) Notify(notifiee Notifiee) {
	n.n.Notify(n)
}

// StopNotify stops signaling events to notifiee.
func (n *Notifier) StopNotify(notifiee Notifiee) {
	n.n.StopNotify(notifiee)
}

// Notifiee is an interface objects must implement to listen to NAT events.
type Notifiee interface {

	// Called every time a successful mapping happens
	// Warning: the port mapping may have changed. If that is the
	// case, both MappingSuccess and MappingChanged are called.
	MappingSuccess(nat *NAT, m Mapping)

	// Called when mapping a port succeeds, but the mapping is
	// with a different port than an earlier success.
	MappingChanged(nat *NAT, m Mapping, oldport int, newport int)

	// Called when a port mapping fails. NAT will continue attempting after
	// the next period. To stop trying, use: mapping.Close(). After this failure,
	// mapping.ExternalPort() will be zero, and nat.ExternalAddrs() will not
	// return the address for this mapping. With luck, the next attempt will
	// succeed, without the client needing to do anything.
	MappingFailed(nat *NAT, m Mapping, oldport int, err error)
}

// Mapping represents a port mapping in a NAT.
type Mapping interface {
	// NAT returns the NAT object this Mapping belongs to.
	NAT() *NAT

	// Protocol returns the protocol of this port mapping. This is either
	// "tcp" or "udp" as no other protocols are likely to be NAT-supported.
	Protocol() string

	// InternalPort returns the internal device port. Mapping will continue to
	// try to map InternalPort() to an external facing port.
	InternalPort() int

	// ExternalPort returns the external facing port. If the mapping is not
	// established, port will be 0
	ExternalPort() int

	// InternalAddr returns the internal address.
	InternalAddr() ma.Multiaddr

	// ExternalAddr returns the external facing address. If the mapping is not
	// established, addr will be nil, and and ErrNoMapping will be returned.
	ExternalAddr() (addr ma.Multiaddr, err error)

	// Close closes the port mapping
	Close() error
}

// keeps republishing
type mapping struct {
	sync.Mutex // guards all fields

	nat     *NAT
	proto   string
	intport int
	extport int
	intaddr ma.Multiaddr
	proc    goprocess.Process

	cached    ma.Multiaddr
	cacheTime time.Time
}

func (m *mapping) NAT() *NAT {
	m.Lock()
	defer m.Unlock()
	return m.nat
}

func (m *mapping) Protocol() string {
	m.Lock()
	defer m.Unlock()
	return m.proto
}

func (m *mapping) InternalPort() int {
	m.Lock()
	defer m.Unlock()
	return m.intport
}

func (m *mapping) ExternalPort() int {
	m.Lock()
	defer m.Unlock()
	return m.extport
}

func (m *mapping) setExternalPort(p int) {
	m.Lock()
	defer m.Unlock()
	m.extport = p
}

func (m *mapping) InternalAddr() ma.Multiaddr {
	m.Lock()
	defer m.Unlock()
	return m.intaddr
}

func (m *mapping) ExternalAddr() (ma.Multiaddr, error) {
	if time.Now().Sub(m.cacheTime) < CacheTime {
		return m.cached, nil
	}

	if m.ExternalPort() == 0 { // dont even try right now.
		return nil, ErrNoMapping
	}

	ip, err := m.nat.nat.GetExternalAddress()
	if err != nil {
		return nil, err
	}

	ipmaddr, err := manet.FromIP(ip)
	if err != nil {
		return nil, fmt.Errorf("error parsing ip")
	}

	// call m.ExternalPort again, as mapping may have changed under our feet. (tocttou)
	extport := m.ExternalPort()
	if extport == 0 {
		return nil, ErrNoMapping
	}

	tcp, err := ma.NewMultiaddr(fmt.Sprintf("/%s/%d", m.Protocol(), extport))
	if err != nil {
		return nil, err
	}

	maddr2 := ipmaddr.Encapsulate(tcp)

	m.cached = maddr2
	m.cacheTime = time.Now()
	return maddr2, nil
}

func (m *mapping) Close() error {
	return m.proc.Close()
}

// Mappings returns a slice of all NAT mappings
func (nat *NAT) Mappings() []Mapping {
	nat.mappingmu.Lock()
	maps2 := make([]Mapping, 0, len(nat.mappings))
	for m := range nat.mappings {
		maps2 = append(maps2, m)
	}
	nat.mappingmu.Unlock()
	return maps2
}

func (nat *NAT) addMapping(m *mapping) {
	// make mapping automatically close when nat is closed.
	nat.proc.AddChild(m.proc)

	nat.mappingmu.Lock()
	nat.mappings[m] = struct{}{}
	nat.mappingmu.Unlock()
}

func (nat *NAT) rmMapping(m *mapping) {
	nat.mappingmu.Lock()
	delete(nat.mappings, m)
	nat.mappingmu.Unlock()
}

// NewMapping attemps to construct a mapping on protocol and internal port
// It will also periodically renew the mapping until the returned Mapping
// -- or its parent NAT -- is Closed.
//
// May not succeed, and mappings may change over time;
// NAT devices may not respect our port requests, and even lie.
// Clients should not store the mapped results, but rather always
// poll our object for the latest mappings.
func (nat *NAT) NewMapping(maddr ma.Multiaddr) (Mapping, error) {
	if nat == nil {
		return nil, fmt.Errorf("no nat available")
	}

	network, addr, err := manet.DialArgs(maddr)
	if err != nil {
		return nil, fmt.Errorf("DialArgs failed on addr:", maddr.String())
	}

	switch network {
	case "tcp", "tcp4", "tcp6":
		network = "tcp"
	case "udp", "udp4", "udp6":
		network = "udp"
	default:
		return nil, fmt.Errorf("transport not supported by NAT: %s", network)
	}

	intports := strings.Split(addr, ":")[1]
	intport, err := strconv.Atoi(intports)
	if err != nil {
		return nil, err
	}

	m := &mapping{
		nat:     nat,
		proto:   network,
		intport: intport,
		intaddr: maddr,
	}
	m.proc = goprocess.WithTeardown(func() error {
		nat.rmMapping(m)
		return nil
	})
	nat.addMapping(m)

	m.proc.AddChild(periodic.Every(MappingDuration/3, func(worker goprocess.Process) {
		nat.establishMapping(m)
	}))

	// do it once synchronously, so first mapping is done right away, and before exiting,
	// allowing users -- in the optimistic case -- to use results right after.
	nat.establishMapping(m)
	return m, nil
}

func (nat *NAT) establishMapping(m *mapping) {
	oldport := m.ExternalPort()
	log.Debugf("Attempting port map: %s/%d", m.Protocol(), m.InternalPort())
	newport, err := nat.nat.AddPortMapping(m.Protocol(), m.InternalPort(), "http", MappingDuration)

	failure := func() {
		m.setExternalPort(0) // clear mapping
		// TODO: log.Event
		log.Debugf("failed to establish port mapping: %s", err)
		nat.Notifier.notifyAll(func(n Notifiee) {
			n.MappingFailed(nat, m, oldport, err)
		})

		// we do not close if the mapping failed,
		// because it may work again next time.
	}

	if err != nil || newport == 0 {
		failure()
		return
	}

	m.setExternalPort(newport)
	ext, err := m.ExternalAddr()
	if err != nil {
		log.Debugf("NAT Mapping addr error: %s %s", m.InternalAddr(), err)
		failure()
		return
	}

	log.Debugf("NAT Mapping: %s --> %s", m.InternalAddr(), ext)
	if oldport != 0 && newport != oldport {
		log.Debugf("failed to renew same port mapping: ch %d -> %d", oldport, newport)
		nat.Notifier.notifyAll(func(n Notifiee) {
			n.MappingChanged(nat, m, oldport, newport)
		})
	}

	nat.Notifier.notifyAll(func(n Notifiee) {
		n.MappingSuccess(nat, m)
	})
}

// PortMapAddrs attempts to open (and continue to keep open)
// port mappings for given addrs. This function blocks until
// all  addresses have been tried. This allows clients to
// retrieve results immediately after:
//
//   nat.PortMapAddrs(addrs)
//   mapped := nat.ExternalAddrs()
//
// Some may not succeed, and mappings may change over time;
// NAT devices may not respect our port requests, and even lie.
// Clients should not store the mapped results, but rather always
// poll our object for the latest mappings.
func (nat *NAT) PortMapAddrs(addrs []ma.Multiaddr) {
	// spin off addr mappings independently.
	var wg sync.WaitGroup
	for _, addr := range addrs {
		// do all of them concurrently
		wg.Add(1)
		go func() {
			defer wg.Done()
			nat.NewMapping(addr)
		}()
	}
	wg.Wait()
}

// MappedAddrs returns address mappings NAT believes have been
// successfully established. Unsuccessful mappings are nil. This is:
//
// 		map[internalAddr]externalAddr
//
// This set of mappings _may not_ be correct, as NAT devices are finicky.
// Consider this with _best effort_ semantics.
func (nat *NAT) MappedAddrs() map[ma.Multiaddr]ma.Multiaddr {

	mappings := nat.Mappings()
	addrmap := make(map[ma.Multiaddr]ma.Multiaddr, len(mappings))

	for _, m := range mappings {
		i := m.InternalAddr()
		e, err := m.ExternalAddr()
		if err != nil {
			addrmap[i] = nil
		} else {
			addrmap[i] = e
		}
	}
	return addrmap
}

// ExternalAddrs returns a list of addresses that NAT believes have
// been successfully established. Unsuccessful mappings are omitted,
// so nat.ExternalAddrs() may return less addresses than nat.InternalAddrs().
// To see which addresses are mapped, use nat.MappedAddrs().
//
// This set of mappings _may not_ be correct, as NAT devices are finicky.
// Consider this with _best effort_ semantics.
func (nat *NAT) ExternalAddrs() []ma.Multiaddr {
	mappings := nat.Mappings()
	addrs := make([]ma.Multiaddr, 0, len(mappings))
	for _, m := range mappings {
		a, err := m.ExternalAddr()
		if err != nil {
			continue // this mapping not currently successful.
		}
		addrs = append(addrs, a)
	}
	return addrs
}
<commit_msg>remove multiple multihash deps<commit_after>package nat

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"sync"
	"time"

	manet "QmRCPT5WRph8aWXmaT2Rfn6ac98YRUUJnNURpD3hNAWp4f/go-multiaddr-net"
	ma "QmbWxL1aXQhBjc1XGjGF1f2KGBMCBYSuT2ThA8YXnXJK83/go-multiaddr"

	goprocess "QmSir6qPL1tjuxd8LkR8VZq6v625ExAUVs2eCLeqQuaPGU/goprocess"
	periodic "QmSir6qPL1tjuxd8LkR8VZq6v625ExAUVs2eCLeqQuaPGU/goprocess/periodic"
	notifier "QmUtEiB6DmXs7eLJiwS9YFyTAtptqzaWutxCsjHy7UKEgo/go-notifier"
	logging "QmWRypnfEwrgH4k93KEHN5hng7VjKYkWmzDYRuTZeh2Mgh/go-log"
	nat "github.com/fd/go-nat"
)

var (
	// ErrNoMapping signals no mapping exists for an address
	ErrNoMapping = errors.New("mapping not established")
)

var log = logging.Logger("nat")

// MappingDuration is a default port mapping duration.
// Port mappings are renewed every (MappingDuration / 3)
const MappingDuration = time.Second * 60

// CacheTime is the time a mapping will cache an external address for
const CacheTime = time.Second * 15

// DiscoverNAT looks for a NAT device in the network and
// returns an object that can manage port mappings.
func DiscoverNAT() *NAT {
	nat, err := nat.DiscoverGateway()
	if err != nil {
		log.Debug("DiscoverGateway error:", err)
		return nil
	}
	addr, err := nat.GetDeviceAddress()
	if err != nil {
		log.Debug("DiscoverGateway address error:", err)
	} else {
		log.Debug("DiscoverGateway address:", addr)
	}
	return newNAT(nat)
}

// NAT is an object that manages address port mappings in
// NATs (Network Address Translators). It is a long-running
// service that will periodically renew port mappings,
// and keep an up-to-date list of all the external addresses.
type NAT struct {
	nat  nat.NAT
	proc goprocess.Process // manages nat mappings lifecycle

	mappingmu sync.RWMutex // guards mappings
	mappings  map[*mapping]struct{}

	Notifier
}

func newNAT(realNAT nat.NAT) *NAT {
	return &NAT{
		nat:      realNAT,
		proc:     goprocess.WithParent(goprocess.Background()),
		mappings: make(map[*mapping]struct{}),
	}
}

// Close shuts down all port mappings. NAT can no longer be used.
func (nat *NAT) Close() error {
	return nat.proc.Close()
}

// Process returns the nat's life-cycle manager, for making it listen
// to close signals.
func (nat *NAT) Process() goprocess.Process {
	return nat.proc
}

// Notifier is an object that assists NAT in notifying listeners.
// It is implemented using thirdparty/notifier
type Notifier struct {
	n notifier.Notifier
}

func (n *Notifier) notifyAll(notify func(n Notifiee)) {
	n.n.NotifyAll(func(n notifier.Notifiee) {
		notify(n.(Notifiee))
	})
}

// Notify signs up notifiee to listen to NAT events.
func (n *Notifier) Notify(notifiee Notifiee) {
	n.n.Notify(n)
}

// StopNotify stops signaling events to notifiee.
func (n *Notifier) StopNotify(notifiee Notifiee) {
	n.n.StopNotify(notifiee)
}

// Notifiee is an interface objects must implement to listen to NAT events.
type Notifiee interface {

	// Called every time a successful mapping happens
	// Warning: the port mapping may have changed. If that is the
	// case, both MappingSuccess and MappingChanged are called.
	MappingSuccess(nat *NAT, m Mapping)

	// Called when mapping a port succeeds, but the mapping is
	// with a different port than an earlier success.
	MappingChanged(nat *NAT, m Mapping, oldport int, newport int)

	// Called when a port mapping fails. NAT will continue attempting after
	// the next period. To stop trying, use: mapping.Close(). After this failure,
	// mapping.ExternalPort() will be zero, and nat.ExternalAddrs() will not
	// return the address for this mapping. With luck, the next attempt will
	// succeed, without the client needing to do anything.
	MappingFailed(nat *NAT, m Mapping, oldport int, err error)
}

// Mapping represents a port mapping in a NAT.
type Mapping interface {
	// NAT returns the NAT object this Mapping belongs to.
	NAT() *NAT

	// Protocol returns the protocol of this port mapping. This is either
	// "tcp" or "udp" as no other protocols are likely to be NAT-supported.
	Protocol() string

	// InternalPort returns the internal device port. Mapping will continue to
	// try to map InternalPort() to an external facing port.
	InternalPort() int

	// ExternalPort returns the external facing port. If the mapping is not
	// established, port will be 0
	ExternalPort() int

	// InternalAddr returns the internal address.
	InternalAddr() ma.Multiaddr

	// ExternalAddr returns the external facing address. If the mapping is not
	// established, addr will be nil, and and ErrNoMapping will be returned.
	ExternalAddr() (addr ma.Multiaddr, err error)

	// Close closes the port mapping
	Close() error
}

// keeps republishing
type mapping struct {
	sync.Mutex // guards all fields

	nat     *NAT
	proto   string
	intport int
	extport int
	intaddr ma.Multiaddr
	proc    goprocess.Process

	cached    ma.Multiaddr
	cacheTime time.Time
}

func (m *mapping) NAT() *NAT {
	m.Lock()
	defer m.Unlock()
	return m.nat
}

func (m *mapping) Protocol() string {
	m.Lock()
	defer m.Unlock()
	return m.proto
}

func (m *mapping) InternalPort() int {
	m.Lock()
	defer m.Unlock()
	return m.intport
}

func (m *mapping) ExternalPort() int {
	m.Lock()
	defer m.Unlock()
	return m.extport
}

func (m *mapping) setExternalPort(p int) {
	m.Lock()
	defer m.Unlock()
	m.extport = p
}

func (m *mapping) InternalAddr() ma.Multiaddr {
	m.Lock()
	defer m.Unlock()
	return m.intaddr
}

func (m *mapping) ExternalAddr() (ma.Multiaddr, error) {
	if time.Now().Sub(m.cacheTime) < CacheTime {
		return m.cached, nil
	}

	if m.ExternalPort() == 0 { // dont even try right now.
		return nil, ErrNoMapping
	}

	ip, err := m.nat.nat.GetExternalAddress()
	if err != nil {
		return nil, err
	}

	ipmaddr, err := manet.FromIP(ip)
	if err != nil {
		return nil, fmt.Errorf("error parsing ip")
	}

	// call m.ExternalPort again, as mapping may have changed under our feet. (tocttou)
	extport := m.ExternalPort()
	if extport == 0 {
		return nil, ErrNoMapping
	}

	tcp, err := ma.NewMultiaddr(fmt.Sprintf("/%s/%d", m.Protocol(), extport))
	if err != nil {
		return nil, err
	}

	maddr2 := ipmaddr.Encapsulate(tcp)

	m.cached = maddr2
	m.cacheTime = time.Now()
	return maddr2, nil
}

func (m *mapping) Close() error {
	return m.proc.Close()
}

// Mappings returns a slice of all NAT mappings
func (nat *NAT) Mappings() []Mapping {
	nat.mappingmu.Lock()
	maps2 := make([]Mapping, 0, len(nat.mappings))
	for m := range nat.mappings {
		maps2 = append(maps2, m)
	}
	nat.mappingmu.Unlock()
	return maps2
}

func (nat *NAT) addMapping(m *mapping) {
	// make mapping automatically close when nat is closed.
	nat.proc.AddChild(m.proc)

	nat.mappingmu.Lock()
	nat.mappings[m] = struct{}{}
	nat.mappingmu.Unlock()
}

func (nat *NAT) rmMapping(m *mapping) {
	nat.mappingmu.Lock()
	delete(nat.mappings, m)
	nat.mappingmu.Unlock()
}

// NewMapping attemps to construct a mapping on protocol and internal port
// It will also periodically renew the mapping until the returned Mapping
// -- or its parent NAT -- is Closed.
//
// May not succeed, and mappings may change over time;
// NAT devices may not respect our port requests, and even lie.
// Clients should not store the mapped results, but rather always
// poll our object for the latest mappings.
func (nat *NAT) NewMapping(maddr ma.Multiaddr) (Mapping, error) {
	if nat == nil {
		return nil, fmt.Errorf("no nat available")
	}

	network, addr, err := manet.DialArgs(maddr)
	if err != nil {
		return nil, fmt.Errorf("DialArgs failed on addr:", maddr.String())
	}

	switch network {
	case "tcp", "tcp4", "tcp6":
		network = "tcp"
	case "udp", "udp4", "udp6":
		network = "udp"
	default:
		return nil, fmt.Errorf("transport not supported by NAT: %s", network)
	}

	intports := strings.Split(addr, ":")[1]
	intport, err := strconv.Atoi(intports)
	if err != nil {
		return nil, err
	}

	m := &mapping{
		nat:     nat,
		proto:   network,
		intport: intport,
		intaddr: maddr,
	}
	m.proc = goprocess.WithTeardown(func() error {
		nat.rmMapping(m)
		return nil
	})
	nat.addMapping(m)

	m.proc.AddChild(periodic.Every(MappingDuration/3, func(worker goprocess.Process) {
		nat.establishMapping(m)
	}))

	// do it once synchronously, so first mapping is done right away, and before exiting,
	// allowing users -- in the optimistic case -- to use results right after.
	nat.establishMapping(m)
	return m, nil
}

func (nat *NAT) establishMapping(m *mapping) {
	oldport := m.ExternalPort()
	log.Debugf("Attempting port map: %s/%d", m.Protocol(), m.InternalPort())
	newport, err := nat.nat.AddPortMapping(m.Protocol(), m.InternalPort(), "http", MappingDuration)

	failure := func() {
		m.setExternalPort(0) // clear mapping
		// TODO: log.Event
		log.Debugf("failed to establish port mapping: %s", err)
		nat.Notifier.notifyAll(func(n Notifiee) {
			n.MappingFailed(nat, m, oldport, err)
		})

		// we do not close if the mapping failed,
		// because it may work again next time.
	}

	if err != nil || newport == 0 {
		failure()
		return
	}

	m.setExternalPort(newport)
	ext, err := m.ExternalAddr()
	if err != nil {
		log.Debugf("NAT Mapping addr error: %s %s", m.InternalAddr(), err)
		failure()
		return
	}

	log.Debugf("NAT Mapping: %s --> %s", m.InternalAddr(), ext)
	if oldport != 0 && newport != oldport {
		log.Debugf("failed to renew same port mapping: ch %d -> %d", oldport, newport)
		nat.Notifier.notifyAll(func(n Notifiee) {
			n.MappingChanged(nat, m, oldport, newport)
		})
	}

	nat.Notifier.notifyAll(func(n Notifiee) {
		n.MappingSuccess(nat, m)
	})
}

// PortMapAddrs attempts to open (and continue to keep open)
// port mappings for given addrs. This function blocks until
// all  addresses have been tried. This allows clients to
// retrieve results immediately after:
//
//   nat.PortMapAddrs(addrs)
//   mapped := nat.ExternalAddrs()
//
// Some may not succeed, and mappings may change over time;
// NAT devices may not respect our port requests, and even lie.
// Clients should not store the mapped results, but rather always
// poll our object for the latest mappings.
func (nat *NAT) PortMapAddrs(addrs []ma.Multiaddr) {
	// spin off addr mappings independently.
	var wg sync.WaitGroup
	for _, addr := range addrs {
		// do all of them concurrently
		wg.Add(1)
		go func() {
			defer wg.Done()
			nat.NewMapping(addr)
		}()
	}
	wg.Wait()
}

// MappedAddrs returns address mappings NAT believes have been
// successfully established. Unsuccessful mappings are nil. This is:
//
// 		map[internalAddr]externalAddr
//
// This set of mappings _may not_ be correct, as NAT devices are finicky.
// Consider this with _best effort_ semantics.
func (nat *NAT) MappedAddrs() map[ma.Multiaddr]ma.Multiaddr {

	mappings := nat.Mappings()
	addrmap := make(map[ma.Multiaddr]ma.Multiaddr, len(mappings))

	for _, m := range mappings {
		i := m.InternalAddr()
		e, err := m.ExternalAddr()
		if err != nil {
			addrmap[i] = nil
		} else {
			addrmap[i] = e
		}
	}
	return addrmap
}

// ExternalAddrs returns a list of addresses that NAT believes have
// been successfully established. Unsuccessful mappings are omitted,
// so nat.ExternalAddrs() may return less addresses than nat.InternalAddrs().
// To see which addresses are mapped, use nat.MappedAddrs().
//
// This set of mappings _may not_ be correct, as NAT devices are finicky.
// Consider this with _best effort_ semantics.
func (nat *NAT) ExternalAddrs() []ma.Multiaddr {
	mappings := nat.Mappings()
	addrs := make([]ma.Multiaddr, 0, len(mappings))
	for _, m := range mappings {
		a, err := m.ExternalAddr()
		if err != nil {
			continue // this mapping not currently successful.
		}
		addrs = append(addrs, a)
	}
	return addrs
}
<|endoftext|>