<commit_before>package main

import (
	"fmt"
	"os/exec"
	"time"

	pbb "github.com/brotherlogic/buildserver/proto"
	pbfc "github.com/brotherlogic/filecopier/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	"golang.org/x/net/context"
)

const (
	pendWait = time.Minute
)

func (s *Server) runTransition(ctx context.Context, job *pb.JobAssignment) {
	startState := job.State
	switch job.State {
	case pb.State_ACKNOWLEDGED:
		key := s.scheduleBuild(ctx, job)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("SCHED: %v @ %v", key, time.Now())
		s.stateMutex.Unlock()
		if !job.Job.Bootstrap {
			if key != "" {
				job.Server = s.Registry.Identifier
				job.State = pb.State_BUILT
				job.RunningVersion = key
			} else {
				// Bootstrap this job since we don't have an initial version
				if job.Job.PartialBootstrap {
					job.Job.Bootstrap = true
				}
			}
		} else {
			job.CommandKey = key
			job.State = pb.State_BUILDING
			job.Server = s.Registry.Identifier
		}
	case pb.State_BUILDING:
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("BUILD(%v): %v", job.CommandKey, s.scheduler.getState(job.CommandKey))
		s.stateMutex.Unlock()
		if s.taskComplete(job.CommandKey) {
			job.State = pb.State_BUILT
		}
	case pb.State_BUILT:
		output, _ := s.scheduler.getOutput(job.CommandKey)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("BUILT(%v): (%v): %v", job.CommandKey, len(output), output)
		s.stateMutex.Unlock()
		if job.Job.Bootstrap && len(output) > 0 {
			if job.BuildFail == 5 {
				s.deliverCrashReport(ctx, job, output)
				job.BuildFail = 0
			}
			job.BuildFail++
			job.State = pb.State_DIED
		} else {
			job.BuildFail = 0
			key := s.scheduleRun(job.Job)
			job.CommandKey = key
			job.StartTime = time.Now().Unix()
			job.State = pb.State_PENDING
			if _, ok := s.pendingMap[time.Now().Weekday()]; !ok {
				s.pendingMap[time.Now().Weekday()] = make(map[string]int)
			}
			s.pendingMap[time.Now().Weekday()][job.Job.Name]++
		}
	case pb.State_PENDING:
		s.stateMutex.Lock()
		out, _ := s.scheduler.getOutput(job.CommandKey)
		s.stateMap[job.Job.Name] = fmt.Sprintf("OUTPUT = %v", out)
		s.stateMutex.Unlock()
		if time.Now().Add(-time.Minute).Unix() > job.StartTime {
			job.State = pb.State_RUNNING
		}
	case pb.State_RUNNING:
		output, errout := s.scheduler.getOutput(job.CommandKey)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("ROUTPUT = %v, %v", output, s.scheduler.getStatus(job.CommandKey))
		job.Status = s.scheduler.getStatus(job.CommandKey)
		s.stateMutex.Unlock()
		if len(job.CommandKey) > 0 && s.taskComplete(job.CommandKey) {
			s.stateMutex.Lock()
			s.stateMap[job.Job.Name] = fmt.Sprintf("COMPLETE = (%v, %v)", job, output)
			s.stateMutex.Unlock()
			s.deliverCrashReport(ctx, job, output)
			job.State = pb.State_DIED
		}

		err := s.discover.discover(job.Job.Name, s.Registry.Identifier)
		if err != nil {
			if job.DiscoverCount > 30 {
				output2, errout2 := s.scheduler.getErrOutput(job.CommandKey)
				s.RaiseIssue(ctx, "Cannot Discover Running Server", fmt.Sprintf("%v on %v is not discoverable, despite running (%v) the output says %v (%v), %v, %v", job.Job.Name, s.Registry.Identifier, err, output, errout, output2, errout2), false)
			}
			job.DiscoverCount++
		} else {
			job.DiscoverCount = 0
		}

		// Restart this job if we need to
		if !job.Job.Bootstrap {
			if time.Now().Sub(time.Unix(job.LastVersionPull, 0)) > time.Minute*5 {
				version, err := s.getVersion(ctx, job.Job)
				job.LastVersionPull = time.Now().Unix()

				if err == nil && version.Version != job.RunningVersion {
					s.stateMutex.Lock()
					s.stateMap[job.Job.Name] = fmt.Sprintf("VERSION_MISMATCH = %v,%v", version, job.RunningVersion)
					s.stateMutex.Unlock()
					s.scheduler.killJob(job.CommandKey)
				}
			}
		}
	case pb.State_BRINK_OF_DEATH:
		if s.version.confirm(ctx, job.Job.Name) {
			s.scheduler.killJob(job.CommandKey)
			job.State = pb.State_ACKNOWLEDGED
		}
	case pb.State_DIED:
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("DIED %v", job.CommandKey)
		s.stateMutex.Unlock()
		s.scheduler.removeJob(job.CommandKey)
		job.State = pb.State_ACKNOWLEDGED
	}

	if job.State != startState {
		job.LastTransitionTime = time.Now().Unix()
	}
}

type translator interface {
	build(job *pb.Job) *exec.Cmd
	run(job *pb.Job) *exec.Cmd
}

type checker interface {
	isAlive(ctx context.Context, job *pb.JobAssignment) bool
}

func (s *Server) getVersion(ctx context.Context, job *pb.Job) (*pbb.Version, error) {
	version, err := s.builder.build(ctx, job)
	if err != nil {
		return &pbb.Version{}, err
	}

	return version, nil

}

func updateJob(err error, job *pb.JobAssignment, resp *pbfc.CopyResponse) {
	if err == nil {
		job.QueuePos = resp.IndexInQueue
	}

}

// scheduleBuild builds out the job, returning the current version
func (s *Server) scheduleBuild(ctx context.Context, job *pb.JobAssignment) string {
	if job.Job.Bootstrap {
		c := s.translator.build(job.Job)
		return s.scheduler.Schedule(&rCommand{command: c, base: job.Job.Name})
	}

	val, err := s.builder.build(ctx, job.Job)
	if err != nil {
		s.Log(fmt.Sprintf("Error on build: %v", err))
		return ""
	}
	return val.Version
}

func (s *Server) scheduleRun(job *pb.Job) string {
	c := s.translator.run(job)
	return s.scheduler.Schedule(&rCommand{command: c, base: job.Name})
}

func (s *Server) taskComplete(key string) bool {
	return s.scheduler.schedulerComplete(key)
}
<commit_msg>Resets bootstrap setting<commit_after>package main

import (
	"fmt"
	"os/exec"
	"time"

	pbb "github.com/brotherlogic/buildserver/proto"
	pbfc "github.com/brotherlogic/filecopier/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	"golang.org/x/net/context"
)

const (
	pendWait = time.Minute
)

func (s *Server) runTransition(ctx context.Context, job *pb.JobAssignment) {
	startState := job.State
	switch job.State {
	case pb.State_ACKNOWLEDGED:
		key := s.scheduleBuild(ctx, job)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("SCHED: %v @ %v", key, time.Now())
		s.stateMutex.Unlock()
		if !job.Job.Bootstrap {
			if key != "" {
				job.Server = s.Registry.Identifier
				job.State = pb.State_BUILT
				job.RunningVersion = key
			} else {
				// Bootstrap this job since we don't have an initial version
				if job.Job.PartialBootstrap {
					job.Job.Bootstrap = true
				}
			}
		} else {
			job.CommandKey = key
			job.State = pb.State_BUILDING
			job.Server = s.Registry.Identifier
		}
	case pb.State_BUILDING:
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("BUILD(%v): %v", job.CommandKey, s.scheduler.getState(job.CommandKey))
		s.stateMutex.Unlock()
		if s.taskComplete(job.CommandKey) {
			job.State = pb.State_BUILT
		}
	case pb.State_BUILT:
		output, _ := s.scheduler.getOutput(job.CommandKey)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("BUILT(%v): (%v): %v", job.CommandKey, len(output), output)
		s.stateMutex.Unlock()
		if job.Job.Bootstrap && len(output) > 0 {
			if job.BuildFail == 5 {
				s.deliverCrashReport(ctx, job, output)
				job.BuildFail = 0
			}
			job.BuildFail++
			job.State = pb.State_DIED
		} else {
			job.BuildFail = 0
			key := s.scheduleRun(job.Job)
			job.CommandKey = key
			job.StartTime = time.Now().Unix()
			job.State = pb.State_PENDING
			if _, ok := s.pendingMap[time.Now().Weekday()]; !ok {
				s.pendingMap[time.Now().Weekday()] = make(map[string]int)
			}
			s.pendingMap[time.Now().Weekday()][job.Job.Name]++
		}
	case pb.State_PENDING:
		if job.Job.PartialBootstrap && job.Job.Bootstrap {
			job.Job.Bootstrap = false
		}
		s.stateMutex.Lock()
		out, _ := s.scheduler.getOutput(job.CommandKey)
		s.stateMap[job.Job.Name] = fmt.Sprintf("OUTPUT = %v", out)
		s.stateMutex.Unlock()
		if time.Now().Add(-time.Minute).Unix() > job.StartTime {
			job.State = pb.State_RUNNING
		}
	case pb.State_RUNNING:
		output, errout := s.scheduler.getOutput(job.CommandKey)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("ROUTPUT = %v, %v", output, s.scheduler.getStatus(job.CommandKey))
		job.Status = s.scheduler.getStatus(job.CommandKey)
		s.stateMutex.Unlock()
		if len(job.CommandKey) > 0 && s.taskComplete(job.CommandKey) {
			s.stateMutex.Lock()
			s.stateMap[job.Job.Name] = fmt.Sprintf("COMPLETE = (%v, %v)", job, output)
			s.stateMutex.Unlock()
			s.deliverCrashReport(ctx, job, output)
			job.State = pb.State_DIED
		}

		err := s.discover.discover(job.Job.Name, s.Registry.Identifier)
		if err != nil {
			if job.DiscoverCount > 30 {
				output2, errout2 := s.scheduler.getErrOutput(job.CommandKey)
				s.RaiseIssue(ctx, "Cannot Discover Running Server", fmt.Sprintf("%v on %v is not discoverable, despite running (%v) the output says %v (%v), %v, %v", job.Job.Name, s.Registry.Identifier, err, output, errout, output2, errout2), false)
			}
			job.DiscoverCount++
		} else {
			job.DiscoverCount = 0
		}

		// Restart this job if we need to
		if !job.Job.Bootstrap {
			if time.Now().Sub(time.Unix(job.LastVersionPull, 0)) > time.Minute*5 {
				version, err := s.getVersion(ctx, job.Job)
				job.LastVersionPull = time.Now().Unix()

				if err == nil && version.Version != job.RunningVersion {
					s.stateMutex.Lock()
					s.stateMap[job.Job.Name] = fmt.Sprintf("VERSION_MISMATCH = %v,%v", version, job.RunningVersion)
					s.stateMutex.Unlock()
					s.scheduler.killJob(job.CommandKey)
				}
			}
		}
	case pb.State_BRINK_OF_DEATH:
		if s.version.confirm(ctx, job.Job.Name) {
			s.scheduler.killJob(job.CommandKey)
			job.State = pb.State_ACKNOWLEDGED
		}
	case pb.State_DIED:
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("DIED %v", job.CommandKey)
		s.stateMutex.Unlock()
		s.scheduler.removeJob(job.CommandKey)
		job.State = pb.State_ACKNOWLEDGED
	}

	if job.State != startState {
		job.LastTransitionTime = time.Now().Unix()
	}
}

type translator interface {
	build(job *pb.Job) *exec.Cmd
	run(job *pb.Job) *exec.Cmd
}

type checker interface {
	isAlive(ctx context.Context, job *pb.JobAssignment) bool
}

func (s *Server) getVersion(ctx context.Context, job *pb.Job) (*pbb.Version, error) {
	version, err := s.builder.build(ctx, job)
	if err != nil {
		return &pbb.Version{}, err
	}

	return version, nil

}

func updateJob(err error, job *pb.JobAssignment, resp *pbfc.CopyResponse) {
	if err == nil {
		job.QueuePos = resp.IndexInQueue
	}

}

// scheduleBuild builds out the job, returning the current version
func (s *Server) scheduleBuild(ctx context.Context, job *pb.JobAssignment) string {
	if job.Job.Bootstrap {
		c := s.translator.build(job.Job)
		return s.scheduler.Schedule(&rCommand{command: c, base: job.Job.Name})
	}

	val, err := s.builder.build(ctx, job.Job)
	if err != nil {
		s.Log(fmt.Sprintf("Error on build: %v", err))
		return ""
	}
	return val.Version
}

func (s *Server) scheduleRun(job *pb.Job) string {
	c := s.translator.run(job)
	return s.scheduler.Schedule(&rCommand{command: c, base: job.Name})
}

func (s *Server) taskComplete(key string) bool {
	return s.scheduler.schedulerComplete(key)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"time"

	pbb "github.com/brotherlogic/buildserver/proto"
	pbfc "github.com/brotherlogic/filecopier/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
)

const (
	pendWait = time.Minute
)

func (s *Server) runTransition(ctx context.Context, job *pb.JobAssignment) {
	startState := job.State
	switch job.State {
	case pb.State_ACKNOWLEDGED:
		key := s.scheduleBuild(ctx, job)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("SCHED: %v @ %v", key, time.Now())
		s.stateMutex.Unlock()
		if !job.Job.Bootstrap {
			if key != "" {
				job.Server = s.Registry.Identifier
				job.State = pb.State_BUILT
				job.RunningVersion = key
			}
		} else {
			job.CommandKey = key
			job.State = pb.State_BUILDING
			job.Server = s.Registry.Identifier
		}
	case pb.State_BUILDING:
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("BUILD(%v): %v", job.CommandKey, s.scheduler.getState(job.CommandKey))
		s.stateMutex.Unlock()
		if s.taskComplete(job.CommandKey) {
			job.State = pb.State_BUILT
		}
	case pb.State_BUILT:
		output, _ := s.scheduler.getOutput(job.CommandKey)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("BUILT(%v): (%v): %v", job.CommandKey, len(output), output)
		s.stateMutex.Unlock()
		if job.Job.Bootstrap && len(output) > 0 {
			if job.BuildFail == 5 {
				s.deliverCrashReport(ctx, job, output)
				job.BuildFail = 0
			}
			job.BuildFail++
			job.State = pb.State_DIED
		} else {
			job.BuildFail = 0
			key := s.scheduleRun(job.Job)
			job.CommandKey = key
			job.StartTime = time.Now().Unix()
			job.State = pb.State_PENDING
			if _, ok := s.pendingMap[time.Now().Weekday()]; !ok {
				s.pendingMap[time.Now().Weekday()] = make(map[string]int)
			}
			s.pendingMap[time.Now().Weekday()][job.Job.Name]++
		}
	case pb.State_PENDING:
		s.stateMutex.Lock()
		out, _ := s.scheduler.getOutput(job.CommandKey)
		s.stateMap[job.Job.Name] = fmt.Sprintf("OUTPUT = %v", out)
		s.stateMutex.Unlock()
		if time.Now().Add(-time.Minute).Unix() > job.StartTime {
			job.State = pb.State_RUNNING
		}
	case pb.State_RUNNING:
		output, errout := s.scheduler.getOutput(job.CommandKey)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("ROUTPUT = %v, %v", output, s.scheduler.getStatus(job.CommandKey))
		job.Status = s.scheduler.getStatus(job.CommandKey)
		s.stateMutex.Unlock()
		if len(job.CommandKey) > 0 && s.taskComplete(job.CommandKey) {
			s.stateMutex.Lock()
			s.stateMap[job.Job.Name] = fmt.Sprintf("COMPLETE = %v", output)
			s.stateMutex.Unlock()
			s.deliverCrashReport(ctx, job, output)
			job.State = pb.State_DIED
		}

		err := s.discover.discover(job.Job.Name, s.Registry.Identifier)
		if err != nil {
			if job.DiscoverCount > 30 {
				output2, errout2 := s.scheduler.getErrOutput(job.CommandKey)
				s.RaiseIssue(ctx, "Cannot Discover Running Server", fmt.Sprintf("%v on %v is not discoverable, despite running (%v) the output says %v (%v), %v, %v", job.Job.Name, s.Registry.Identifier, err, output, errout, output2, errout2), false)
			}
			job.DiscoverCount++
		} else {
			job.DiscoverCount = 0
		}

		// Restart this job if we need to
		if !job.Job.Bootstrap {
			if time.Now().Sub(time.Unix(job.LastVersionPull, 0)) > time.Minute*5 {
				version, err := s.getVersion(ctx, job.Job)
				job.LastVersionPull = time.Now().Unix()

				if err == nil && version.Version != job.RunningVersion {
					s.stateMutex.Lock()
					s.stateMap[job.Job.Name] = fmt.Sprintf("VERSION_MISMATCH = %v,%v", version, job.RunningVersion)
					s.stateMutex.Unlock()
					s.scheduler.killJob(job.CommandKey)
				}
			}
		}
	case pb.State_BRINK_OF_DEATH:
		if s.version.confirm(ctx, job.Job.Name) {
			s.scheduler.killJob(job.CommandKey)
			job.State = pb.State_ACKNOWLEDGED
		}
	case pb.State_DIED:
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("DIED %v", job.CommandKey)
		s.stateMutex.Unlock()
		s.scheduler.removeJob(job.CommandKey)
		job.State = pb.State_ACKNOWLEDGED
	}

	if job.State != startState {
		job.LastTransitionTime = time.Now().Unix()
	}
}

type translator interface {
	build(job *pb.Job) *exec.Cmd
	run(job *pb.Job) *exec.Cmd
}

type checker interface {
	isAlive(ctx context.Context, job *pb.JobAssignment) bool
}

func (s *Server) getVersion(ctx context.Context, job *pb.Job) (*pbb.Version, error) {
	versions, err := s.builder.build(ctx, job)
	if err != nil {
		return &pbb.Version{}, err
	}

	if len(versions) == 0 {
		return &pbb.Version{}, nil
	}

	return versions[0], nil

}

func updateJob(err error, job *pb.JobAssignment, resp *pbfc.CopyResponse) {
	if err == nil {
		job.QueuePos = resp.IndexInQueue
	}

}

// scheduleBuild builds out the job, returning the current version
func (s *Server) scheduleBuild(ctx context.Context, job *pb.JobAssignment) string {
	if job.Job.Bootstrap {
		c := s.translator.build(job.Job)
		return s.scheduler.Schedule(&rCommand{command: c, base: job.Job.Name})
	}

	versions, err := s.builder.build(ctx, job.Job)

	s.lastCopyStatus = fmt.Sprintf("%v", err)
	if len(versions) == 0 {
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("No Versions: %v", err)
		s.stateMutex.Unlock()
		return ""
	}

	t := time.Now()

	//Only copy if the latest version is different to the local version
	s.versionsMutex.Lock()
	v, ok := s.versions[job.Job.Name]
	s.versionsMutex.Unlock()
	if !ok || v.Version != versions[0].Version {
		s.copies++

		resp, err := s.builder.copy(ctx, versions[0])
		s.lastCopyTime = time.Now().Sub(t)
		s.lastCopyStatus = fmt.Sprintf("%v", err)
		if err != nil || resp.Status != pbfc.CopyStatus_COMPLETE || len(resp.Error) > 0 {
			updateJob(err, job, resp)
			s.stateMutex.Lock()
			s.stateMap[job.Job.Name] = fmt.Sprintf("Copy fail (%v) -> %v", time.Now().Sub(t), err)
			s.stateMutex.Unlock()
			return ""
		}
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("Copied version %v", versions[0].Version)
		s.stateMutex.Unlock()

		//Save the version file alongside the binary
		data, _ := proto.Marshal(versions[0])
		ioutil.WriteFile("/home/simon/gobuild/bin/"+job.Job.Name+".version", data, 0644)
	} else {
		s.skippedCopies++
	}

	s.versionsMutex.Lock()
	s.versions[job.Job.Name] = versions[0]
	s.versionsMutex.Unlock()
	return versions[0].Version
}

func (s *Server) scheduleRun(job *pb.Job) string {
	c := s.translator.run(job)
	return s.scheduler.Schedule(&rCommand{command: c, base: job.Name})
}

func (s *Server) taskComplete(key string) bool {
	return s.scheduler.schedulerComplete(key)
}
<commit_msg>Update<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os/exec"
	"time"

	pbb "github.com/brotherlogic/buildserver/proto"
	pbfc "github.com/brotherlogic/filecopier/proto"
	pb "github.com/brotherlogic/gobuildslave/proto"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
)

const (
	pendWait = time.Minute
)

func (s *Server) runTransition(ctx context.Context, job *pb.JobAssignment) {
	startState := job.State
	switch job.State {
	case pb.State_ACKNOWLEDGED:
		key := s.scheduleBuild(ctx, job)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("SCHED: %v @ %v", key, time.Now())
		s.stateMutex.Unlock()
		if !job.Job.Bootstrap {
			if key != "" {
				job.Server = s.Registry.Identifier
				job.State = pb.State_BUILT
				job.RunningVersion = key
			}
		} else {
			job.CommandKey = key
			job.State = pb.State_BUILDING
			job.Server = s.Registry.Identifier
		}
	case pb.State_BUILDING:
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("BUILD(%v): %v", job.CommandKey, s.scheduler.getState(job.CommandKey))
		s.stateMutex.Unlock()
		if s.taskComplete(job.CommandKey) {
			job.State = pb.State_BUILT
		}
	case pb.State_BUILT:
		output, _ := s.scheduler.getOutput(job.CommandKey)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("BUILT(%v): (%v): %v", job.CommandKey, len(output), output)
		s.stateMutex.Unlock()
		if job.Job.Bootstrap && len(output) > 0 {
			if job.BuildFail == 5 {
				s.deliverCrashReport(ctx, job, output)
				job.BuildFail = 0
			}
			job.BuildFail++
			job.State = pb.State_DIED
		} else {
			job.BuildFail = 0
			key := s.scheduleRun(job.Job)
			job.CommandKey = key
			job.StartTime = time.Now().Unix()
			job.State = pb.State_PENDING
			if _, ok := s.pendingMap[time.Now().Weekday()]; !ok {
				s.pendingMap[time.Now().Weekday()] = make(map[string]int)
			}
			s.pendingMap[time.Now().Weekday()][job.Job.Name]++
		}
	case pb.State_PENDING:
		s.stateMutex.Lock()
		out, _ := s.scheduler.getOutput(job.CommandKey)
		s.stateMap[job.Job.Name] = fmt.Sprintf("OUTPUT = %v", out)
		s.stateMutex.Unlock()
		if time.Now().Add(-time.Minute).Unix() > job.StartTime {
			job.State = pb.State_RUNNING
		}
	case pb.State_RUNNING:
		output, errout := s.scheduler.getOutput(job.CommandKey)
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("ROUTPUT = %v, %v", output, s.scheduler.getStatus(job.CommandKey))
		job.Status = s.scheduler.getStatus(job.CommandKey)
		s.stateMutex.Unlock()
		if len(job.CommandKey) > 0 && s.taskComplete(job.CommandKey) {
			s.stateMutex.Lock()
			s.stateMap[job.Job.Name] = fmt.Sprintf("COMPLETE = %v", output)
			s.stateMutex.Unlock()
			s.deliverCrashReport(ctx, job, output+fmt.Sprintf("%v", errout))
			job.State = pb.State_DIED
		}

		err := s.discover.discover(job.Job.Name, s.Registry.Identifier)
		if err != nil {
			if job.DiscoverCount > 30 {
				output2, errout2 := s.scheduler.getErrOutput(job.CommandKey)
				s.RaiseIssue(ctx, "Cannot Discover Running Server", fmt.Sprintf("%v on %v is not discoverable, despite running (%v) the output says %v (%v), %v, %v", job.Job.Name, s.Registry.Identifier, err, output, errout, output2, errout2), false)
			}
			job.DiscoverCount++
		} else {
			job.DiscoverCount = 0
		}

		// Restart this job if we need to
		if !job.Job.Bootstrap {
			if time.Now().Sub(time.Unix(job.LastVersionPull, 0)) > time.Minute*5 {
				version, err := s.getVersion(ctx, job.Job)
				job.LastVersionPull = time.Now().Unix()

				if err == nil && version.Version != job.RunningVersion {
					s.stateMutex.Lock()
					s.stateMap[job.Job.Name] = fmt.Sprintf("VERSION_MISMATCH = %v,%v", version, job.RunningVersion)
					s.stateMutex.Unlock()
					s.scheduler.killJob(job.CommandKey)
				}
			}
		}
	case pb.State_BRINK_OF_DEATH:
		if s.version.confirm(ctx, job.Job.Name) {
			s.scheduler.killJob(job.CommandKey)
			job.State = pb.State_ACKNOWLEDGED
		}
	case pb.State_DIED:
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("DIED %v", job.CommandKey)
		s.stateMutex.Unlock()
		s.scheduler.removeJob(job.CommandKey)
		job.State = pb.State_ACKNOWLEDGED
	}

	if job.State != startState {
		job.LastTransitionTime = time.Now().Unix()
	}
}

type translator interface {
	build(job *pb.Job) *exec.Cmd
	run(job *pb.Job) *exec.Cmd
}

type checker interface {
	isAlive(ctx context.Context, job *pb.JobAssignment) bool
}

func (s *Server) getVersion(ctx context.Context, job *pb.Job) (*pbb.Version, error) {
	versions, err := s.builder.build(ctx, job)
	if err != nil {
		return &pbb.Version{}, err
	}

	if len(versions) == 0 {
		return &pbb.Version{}, nil
	}

	return versions[0], nil

}

func updateJob(err error, job *pb.JobAssignment, resp *pbfc.CopyResponse) {
	if err == nil {
		job.QueuePos = resp.IndexInQueue
	}

}

// scheduleBuild builds out the job, returning the current version
func (s *Server) scheduleBuild(ctx context.Context, job *pb.JobAssignment) string {
	if job.Job.Bootstrap {
		c := s.translator.build(job.Job)
		return s.scheduler.Schedule(&rCommand{command: c, base: job.Job.Name})
	}

	versions, err := s.builder.build(ctx, job.Job)

	s.lastCopyStatus = fmt.Sprintf("%v", err)
	if len(versions) == 0 {
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("No Versions: %v", err)
		s.stateMutex.Unlock()
		return ""
	}

	t := time.Now()

	//Only copy if the latest version is different to the local version
	s.versionsMutex.Lock()
	v, ok := s.versions[job.Job.Name]
	s.versionsMutex.Unlock()
	if !ok || v.Version != versions[0].Version {
		s.copies++

		resp, err := s.builder.copy(ctx, versions[0])
		s.lastCopyTime = time.Now().Sub(t)
		s.lastCopyStatus = fmt.Sprintf("%v", err)
		if err != nil || resp.Status != pbfc.CopyStatus_COMPLETE || len(resp.Error) > 0 {
			updateJob(err, job, resp)
			s.stateMutex.Lock()
			s.stateMap[job.Job.Name] = fmt.Sprintf("Copy fail (%v) -> %v", time.Now().Sub(t), err)
			s.stateMutex.Unlock()
			return ""
		}
		s.stateMutex.Lock()
		s.stateMap[job.Job.Name] = fmt.Sprintf("Copied version %v", versions[0].Version)
		s.stateMutex.Unlock()

		//Save the version file alongside the binary
		data, _ := proto.Marshal(versions[0])
		ioutil.WriteFile("/home/simon/gobuild/bin/"+job.Job.Name+".version", data, 0644)
	} else {
		s.skippedCopies++
	}

	s.versionsMutex.Lock()
	s.versions[job.Job.Name] = versions[0]
	s.versionsMutex.Unlock()
	return versions[0].Version
}

func (s *Server) scheduleRun(job *pb.Job) string {
	c := s.translator.run(job)
	return s.scheduler.Schedule(&rCommand{command: c, base: job.Name})
}

func (s *Server) taskComplete(key string) bool {
	return s.scheduler.schedulerComplete(key)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/kisielk/gostatsd/statsd"
	"log"
	"time"
)

const (
	defaultMetricsAddr   = ":8125"
	defaultConsoleAddr   = ":8126"
	defaultGraphiteAddr  = "localhost:2003"
	defaultFlushInterval = 10 * time.Second
)

func main() {
	metricsAddr := flag.String("l", defaultMetricsAddr, "address on which to listen for metrics")
	graphiteAddr := flag.String("g", defaultGraphiteAddr, "address of the graphite server")
	flushInterval := flag.Duration("f", defaultFlushInterval, "how often to flush metrics to the graphite server")
	webConsoleAddr := flag.String("web", "", "if set, use as the address of the web-based console")
	consoleAddr := flag.String("console", "", "if set, use as the address of the telnet-based console ")
	flag.Parse()

	// Start the metric aggregator
	graphite, err := statsd.NewGraphiteClient(*graphiteAddr)
	if err != nil {
		log.Fatal(err)
	}
	aggregator := statsd.NewMetricAggregator(&graphite, *flushInterval)
	go aggregator.Aggregate()

	// Start the metric receiver
	f := func(metric statsd.Metric) {
		aggregator.MetricChan <- metric
	}
	receiver := statsd.MetricReceiver{*metricsAddr, statsd.HandlerFunc(f)}
	go receiver.ListenAndReceive()

	// Start the console(s)
	if *consoleAddr != "" {
		console := statsd.ConsoleServer{*consoleAddr, &aggregator}
		go console.ListenAndServe()
	}
	if *webConsoleAddr != "" {
		console := statsd.WebConsoleServer{*webConsoleAddr, &aggregator}
		go console.ListenAndServe()
	}

	// Listen forever
	select {}
}
<commit_msg>modify import from fabware<commit_after>package main

import (
	"flag"
	"github.com/fabware/gostatsd/statsd"
	"log"
	"time"
)

const (
	defaultMetricsAddr   = ":8125"
	defaultConsoleAddr   = ":8126"
	defaultGraphiteAddr  = "localhost:2003"
	defaultFlushInterval = 10 * time.Second
)

func main() {
	metricsAddr := flag.String("l", defaultMetricsAddr, "address on which to listen for metrics")
	graphiteAddr := flag.String("g", defaultGraphiteAddr, "address of the graphite server")
	flushInterval := flag.Duration("f", defaultFlushInterval, "how often to flush metrics to the graphite server")
	webConsoleAddr := flag.String("web", "", "if set, use as the address of the web-based console")
	consoleAddr := flag.String("console", "", "if set, use as the address of the telnet-based console ")
	flag.Parse()

	// Start the metric aggregator
	graphite, err := statsd.NewGraphiteClient(*graphiteAddr)
	if err != nil {
		log.Fatal(err)
	}
	aggregator := statsd.NewMetricAggregator(&graphite, *flushInterval)
	go aggregator.Aggregate()

	// Start the metric receiver
	f := func(metric statsd.Metric) {
		aggregator.MetricChan <- metric
	}
	receiver := statsd.MetricReceiver{*metricsAddr, statsd.HandlerFunc(f)}
	go receiver.ListenAndReceive()

	// Start the console(s)
	if *consoleAddr != "" {
		console := statsd.ConsoleServer{*consoleAddr, &aggregator}
		go console.ListenAndServe()
	}
	if *webConsoleAddr != "" {
		console := statsd.WebConsoleServer{*webConsoleAddr, &aggregator}
		go console.ListenAndServe()
	}

	// Listen forever
	select {}
}
<|endoftext|>


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

package graph

import "github.com/cayleygraph/cayley/quad"

// Transaction stores a bunch of Deltas to apply atomatically on the database.
type Transaction struct {
	// Deltas stores the deltas in the right order
	Deltas []Delta
	// deltas stores the deltas in a map to avoid duplications
	deltas map[Delta]struct{}
}

// NewTransaction initialize a new transaction.
func NewTransaction() *Transaction {
	return &Transaction{Deltas: make([]Delta, 0, 10), deltas: make(map[Delta]struct{}, 10)}
}

// AddQuad adds a new quad to the transaction if it is not already present in it.
// If there is a 'remove' delta for that quad, it will remove that delta from
// the transaction instead of actually addind the quad.
func (t *Transaction) AddQuad(q quad.Quad) {
	ad, rd := createDeltas(q)

	if _, adExists := t.deltas[ad]; !adExists {
		if _, rdExists := t.deltas[rd]; rdExists {
			t.deleteDelta(rd)
		} else {
			t.addDelta(ad)
		}
	}
}

// RemoveQuad adds a quad to remove to the transaction.
// The quad will be removed from the database if it is not present in the
// transaction, otherwise it simply remove it from the transaction.
func (t *Transaction) RemoveQuad(q quad.Quad) {
	ad, rd := createDeltas(q)

	if _, adExists := t.deltas[ad]; adExists {
		t.deleteDelta(ad)
	} else {
		if _, rdExists := t.deltas[rd]; !rdExists {
			t.addDelta(rd)
		}
	}
}

func createDeltas(q quad.Quad) (ad, rd Delta) {
	ad = Delta{
		Quad:   q,
		Action: Add,
	}
	rd = Delta{
		Quad:   q,
		Action: Delete,
	}
	return
}

func (t *Transaction) addDelta(d Delta) {
	t.Deltas = append(t.Deltas, d)
	t.deltas[d] = struct{}{}
}

func (t *Transaction) deleteDelta(d Delta) {
	delete(t.deltas, d)

	for i, id := range t.Deltas {
		if id == d {
			t.Deltas = append(t.Deltas[:i], t.Deltas[i+1:]...)
			break
		}
	}
}
<commit_msg>Update transaction.go<commit_after>// Copyright 2015 The Cayley Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package graph

import "github.com/cayleygraph/cayley/quad"

// Transaction stores a bunch of Deltas to apply automatically on the database.
type Transaction struct {
	// Deltas stores the deltas in the right order
	Deltas []Delta
	// deltas stores the deltas in a map to avoid duplications
	deltas map[Delta]struct{}
}

// NewTransaction initialize a new transaction.
func NewTransaction() *Transaction {
	return &Transaction{Deltas: make([]Delta, 0, 10), deltas: make(map[Delta]struct{}, 10)}
}

// AddQuad adds a new quad to the transaction if it is not already present in it.
// If there is a 'remove' delta for that quad, it will remove that delta from
// the transaction instead of actually adding the quad.
func (t *Transaction) AddQuad(q quad.Quad) {
	ad, rd := createDeltas(q)

	if _, adExists := t.deltas[ad]; !adExists {
		if _, rdExists := t.deltas[rd]; rdExists {
			t.deleteDelta(rd)
		} else {
			t.addDelta(ad)
		}
	}
}

// RemoveQuad adds a quad to remove to the transaction.
// The quad will be removed from the database if it is not present in the
// transaction, otherwise it simply remove it from the transaction.
func (t *Transaction) RemoveQuad(q quad.Quad) {
	ad, rd := createDeltas(q)

	if _, adExists := t.deltas[ad]; adExists {
		t.deleteDelta(ad)
	} else {
		if _, rdExists := t.deltas[rd]; !rdExists {
			t.addDelta(rd)
		}
	}
}

func createDeltas(q quad.Quad) (ad, rd Delta) {
	ad = Delta{
		Quad:   q,
		Action: Add,
	}
	rd = Delta{
		Quad:   q,
		Action: Delete,
	}
	return
}

func (t *Transaction) addDelta(d Delta) {
	t.Deltas = append(t.Deltas, d)
	t.deltas[d] = struct{}{}
}

func (t *Transaction) deleteDelta(d Delta) {
	delete(t.deltas, d)

	for i, id := range t.Deltas {
		if id == d {
			t.Deltas = append(t.Deltas[:i], t.Deltas[i+1:]...)
			break
		}
	}
}
<|endoftext|>


<commit_before>package migration

import (
	"fmt"
	"io"

	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/ioprogress"
	"github.com/lxc/lxd/shared/units"
)

// Type represents the migration transport type. It indicates the method by which the migration can
// take place and what optional features are available.
type Type struct {
	FSType   MigrationFSType // Transport mode selected.
	Features []string        // Feature hints for selected FSType transport mode.
}

// VolumeSourceArgs represents the arguments needed to setup a volume migration source.
type VolumeSourceArgs struct {
	Name          string
	Snapshots     []string
	MigrationType Type
	TrackProgress bool
	MultiSync     bool
	FinalSync     bool
	Data          interface{} // Optional store to persist storage driver state between MultiSync phases.
	ContentType   string
}

// VolumeTargetArgs represents the arguments needed to setup a volume migration sink.
type VolumeTargetArgs struct {
	Name          string
	Description   string
	Config        map[string]string
	Snapshots     []string
	MigrationType Type
	TrackProgress bool
	Refresh       bool
	Live          bool
	VolumeSize    int64
	ContentType   string
}

// TypesToHeader converts one or more Types to a MigrationHeader. It uses the first type argument
// supplied to indicate the preferred migration method and sets the MigrationHeader's Fs type
// to that. If the preferred type is ZFS then it will also set the header's optional ZfsFeatures.
// If the fallback Rsync type is present in any of the types even if it is not preferred, then its
// optional features are added to the header's RsyncFeatures, allowing for fallback negotiation to
// take place on the farside.
func TypesToHeader(types ...Type) *MigrationHeader {
	missingFeature := false
	hasFeature := true
	var preferredType Type

	if len(types) > 0 {
		preferredType = types[0]
	}

	header := MigrationHeader{Fs: &preferredType.FSType}

	// Add ZFS features if preferred type is ZFS.
	if preferredType.FSType == MigrationFSType_ZFS {
		features := ZfsFeatures{
			Compress: &missingFeature,
		}
		for _, feature := range preferredType.Features {
			if feature == "compress" {
				features.Compress = &hasFeature
			}
		}

		header.ZfsFeatures = &features
	}

	// Add BTRFS features if preferred type is BTRFS.
	if preferredType.FSType == MigrationFSType_BTRFS {
		features := BtrfsFeatures{
			MigrationHeader:  &missingFeature,
			HeaderSubvolumes: &missingFeature,
		}
		for _, feature := range preferredType.Features {
			if feature == BTRFSFeatureMigrationHeader {
				features.MigrationHeader = &hasFeature
			} else if feature == BTRFSFeatureSubvolumes {
				features.HeaderSubvolumes = &hasFeature
			}
		}

		header.BtrfsFeatures = &features
	}

	// Check all the types for an Rsync method, if found add its features to the header's RsyncFeatures list.
	for _, t := range types {
		if t.FSType != MigrationFSType_RSYNC && t.FSType != MigrationFSType_BLOCK_AND_RSYNC {
			continue
		}

		features := RsyncFeatures{
			Xattrs:        &missingFeature,
			Delete:        &missingFeature,
			Compress:      &missingFeature,
			Bidirectional: &missingFeature,
		}

		for _, feature := range t.Features {
			if feature == "xattrs" {
				features.Xattrs = &hasFeature
			} else if feature == "delete" {
				features.Delete = &hasFeature
			} else if feature == "compress" {
				features.Compress = &hasFeature
			} else if feature == "bidirectional" {
				features.Bidirectional = &hasFeature
			}
		}

		header.RsyncFeatures = &features
		break // Only use the first rsync transport type found to generate rsync features list.
	}

	return &header
}

// MatchTypes attempts to find matching migration transport types between an offered type sent from a remote
// source and the types supported by a local storage pool. If matches are found then one or more Types are
// returned containing the method and the matching optional features present in both. The function also takes a
// fallback type which is used as an additional offer type preference in case the preferred remote type is not
// compatible with the local type available. It is expected that both sides of the migration will support the
// fallback type for the volume's content type that is being migrated.
func MatchTypes(offer *MigrationHeader, fallbackType MigrationFSType, ourTypes []Type) ([]Type, error) {
	// Generate an offer types slice from the preferred type supplied from remote and the
	// fallback type supplied based on the content type of the transfer.
	offeredFSTypes := []MigrationFSType{offer.GetFs(), fallbackType}

	matchedTypes := []Type{}

	// Find first matching type.
	for _, ourType := range ourTypes {
		for _, offerFSType := range offeredFSTypes {
			if offerFSType != ourType.FSType {
				continue // Not a match, try the next one.
			}

			// We got a match, now extract the relevant offered features.
			var offeredFeatures []string
			if offerFSType == MigrationFSType_ZFS {
				offeredFeatures = offer.GetZfsFeaturesSlice()
			} else if offerFSType == MigrationFSType_BTRFS {
				offeredFeatures = offer.GetBtrfsFeaturesSlice()
			} else if offerFSType == MigrationFSType_RSYNC {
				offeredFeatures = offer.GetRsyncFeaturesSlice()
				if !shared.StringInSlice("bidirectional", offeredFeatures) {
					// If no bi-directional support, this means we are getting a response from
					// an old LXD server that doesn't support bidirectional negotiation, so
					// assume LXD 3.7 level. NOTE: Do NOT extend this list of arguments.
					offeredFeatures = []string{"xattrs", "delete", "compress"}
				}
			}

			// Find common features in both our type and offered type.
			commonFeatures := []string{}
			for _, ourFeature := range ourType.Features {
				if shared.StringInSlice(ourFeature, offeredFeatures) {
					commonFeatures = append(commonFeatures, ourFeature)
				}
			}

			// Append type with combined features.
			matchedTypes = append(matchedTypes, Type{
				FSType:   ourType.FSType,
				Features: commonFeatures,
			})
		}
	}

	if len(matchedTypes) < 1 {
		// No matching transport type found, generate an error with offered types and our types.
		offeredTypeStrings := make([]string, 0, len(offeredFSTypes))
		for _, offerFSType := range offeredFSTypes {
			offeredTypeStrings = append(offeredTypeStrings, offerFSType.String())
		}

		ourTypeStrings := make([]string, 0, len(ourTypes))
		for _, ourType := range ourTypes {
			ourTypeStrings = append(ourTypeStrings, ourType.FSType.String())
		}

		return matchedTypes, fmt.Errorf("No matching migration types found. Offered types: %v, our types: %v", offeredTypeStrings, ourTypeStrings)
	}

	return matchedTypes, nil
}

func progressWrapperRender(op *operations.Operation, key string, description string, progressInt int64, speedInt int64) {
	meta := op.Metadata()
	if meta == nil {
		meta = make(map[string]interface{})
	}

	progress := fmt.Sprintf("%s (%s/s)", units.GetByteSizeString(progressInt, 2), units.GetByteSizeString(speedInt, 2))
	if description != "" {
		progress = fmt.Sprintf("%s: %s (%s/s)", description, units.GetByteSizeString(progressInt, 2), units.GetByteSizeString(speedInt, 2))
	}

	if meta[key] != progress {
		meta[key] = progress
		op.UpdateMetadata(meta)
	}
}

// ProgressReader reports the read progress.
func ProgressReader(op *operations.Operation, key string, description string) func(io.ReadCloser) io.ReadCloser {
	return func(reader io.ReadCloser) io.ReadCloser {
		if op == nil {
			return reader
		}

		progress := func(progressInt int64, speedInt int64) {
			progressWrapperRender(op, key, description, progressInt, speedInt)
		}

		readPipe := &ioprogress.ProgressReader{
			ReadCloser: reader,
			Tracker: &ioprogress.ProgressTracker{
				Handler: progress,
			},
		}

		return readPipe
	}
}

// ProgressWriter reports the write progress.
func ProgressWriter(op *operations.Operation, key string, description string) func(io.WriteCloser) io.WriteCloser {
	return func(writer io.WriteCloser) io.WriteCloser {
		if op == nil {
			return writer
		}

		progress := func(progressInt int64, speedInt int64) {
			progressWrapperRender(op, key, description, progressInt, speedInt)
		}

		writePipe := &ioprogress.ProgressWriter{
			WriteCloser: writer,
			Tracker: &ioprogress.ProgressTracker{
				Handler: progress,
			},
		}

		return writePipe
	}
}

// ProgressTracker returns a migration I/O tracker
func ProgressTracker(op *operations.Operation, key string, description string) *ioprogress.ProgressTracker {
	progress := func(progressInt int64, speedInt int64) {
		progressWrapperRender(op, key, description, progressInt, speedInt)
	}

	tracker := &ioprogress.ProgressTracker{
		Handler: progress,
	}

	return tracker
}
<commit_msg>lxd/migration: Adds AllowInconsistent flag to VolumeSourceArgs.<commit_after>package migration

import (
	"fmt"
	"io"

	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/ioprogress"
	"github.com/lxc/lxd/shared/units"
)

// Type represents the migration transport type. It indicates the method by which the migration can
// take place and what optional features are available.
type Type struct {
	FSType   MigrationFSType // Transport mode selected.
	Features []string        // Feature hints for selected FSType transport mode.
}

// VolumeSourceArgs represents the arguments needed to setup a volume migration source.
type VolumeSourceArgs struct {
	Name              string
	Snapshots         []string
	MigrationType     Type
	TrackProgress     bool
	MultiSync         bool
	FinalSync         bool
	Data              interface{} // Optional store to persist storage driver state between MultiSync phases.
	ContentType       string
	AllowInconsistent bool
}

// VolumeTargetArgs represents the arguments needed to setup a volume migration sink.
type VolumeTargetArgs struct {
	Name          string
	Description   string
	Config        map[string]string
	Snapshots     []string
	MigrationType Type
	TrackProgress bool
	Refresh       bool
	Live          bool
	VolumeSize    int64
	ContentType   string
}

// TypesToHeader converts one or more Types to a MigrationHeader. It uses the first type argument
// supplied to indicate the preferred migration method and sets the MigrationHeader's Fs type
// to that. If the preferred type is ZFS then it will also set the header's optional ZfsFeatures.
// If the fallback Rsync type is present in any of the types even if it is not preferred, then its
// optional features are added to the header's RsyncFeatures, allowing for fallback negotiation to
// take place on the farside.
func TypesToHeader(types ...Type) *MigrationHeader {
	missingFeature := false
	hasFeature := true
	var preferredType Type

	if len(types) > 0 {
		preferredType = types[0]
	}

	header := MigrationHeader{Fs: &preferredType.FSType}

	// Add ZFS features if preferred type is ZFS.
	if preferredType.FSType == MigrationFSType_ZFS {
		features := ZfsFeatures{
			Compress: &missingFeature,
		}
		for _, feature := range preferredType.Features {
			if feature == "compress" {
				features.Compress = &hasFeature
			}
		}

		header.ZfsFeatures = &features
	}

	// Add BTRFS features if preferred type is BTRFS.
	if preferredType.FSType == MigrationFSType_BTRFS {
		features := BtrfsFeatures{
			MigrationHeader:  &missingFeature,
			HeaderSubvolumes: &missingFeature,
		}
		for _, feature := range preferredType.Features {
			if feature == BTRFSFeatureMigrationHeader {
				features.MigrationHeader = &hasFeature
			} else if feature == BTRFSFeatureSubvolumes {
				features.HeaderSubvolumes = &hasFeature
			}
		}

		header.BtrfsFeatures = &features
	}

	// Check all the types for an Rsync method, if found add its features to the header's RsyncFeatures list.
	for _, t := range types {
		if t.FSType != MigrationFSType_RSYNC && t.FSType != MigrationFSType_BLOCK_AND_RSYNC {
			continue
		}

		features := RsyncFeatures{
			Xattrs:        &missingFeature,
			Delete:        &missingFeature,
			Compress:      &missingFeature,
			Bidirectional: &missingFeature,
		}

		for _, feature := range t.Features {
			if feature == "xattrs" {
				features.Xattrs = &hasFeature
			} else if feature == "delete" {
				features.Delete = &hasFeature
			} else if feature == "compress" {
				features.Compress = &hasFeature
			} else if feature == "bidirectional" {
				features.Bidirectional = &hasFeature
			}
		}

		header.RsyncFeatures = &features
		break // Only use the first rsync transport type found to generate rsync features list.
	}

	return &header
}

// MatchTypes attempts to find matching migration transport types between an offered type sent from a remote
// source and the types supported by a local storage pool. If matches are found then one or more Types are
// returned containing the method and the matching optional features present in both. The function also takes a
// fallback type which is used as an additional offer type preference in case the preferred remote type is not
// compatible with the local type available. It is expected that both sides of the migration will support the
// fallback type for the volume's content type that is being migrated.
func MatchTypes(offer *MigrationHeader, fallbackType MigrationFSType, ourTypes []Type) ([]Type, error) {
	// Generate an offer types slice from the preferred type supplied from remote and the
	// fallback type supplied based on the content type of the transfer.
	offeredFSTypes := []MigrationFSType{offer.GetFs(), fallbackType}

	matchedTypes := []Type{}

	// Find first matching type.
	for _, ourType := range ourTypes {
		for _, offerFSType := range offeredFSTypes {
			if offerFSType != ourType.FSType {
				continue // Not a match, try the next one.
			}

			// We got a match, now extract the relevant offered features.
			var offeredFeatures []string
			if offerFSType == MigrationFSType_ZFS {
				offeredFeatures = offer.GetZfsFeaturesSlice()
			} else if offerFSType == MigrationFSType_BTRFS {
				offeredFeatures = offer.GetBtrfsFeaturesSlice()
			} else if offerFSType == MigrationFSType_RSYNC {
				offeredFeatures = offer.GetRsyncFeaturesSlice()
				if !shared.StringInSlice("bidirectional", offeredFeatures) {
					// If no bi-directional support, this means we are getting a response from
					// an old LXD server that doesn't support bidirectional negotiation, so
					// assume LXD 3.7 level. NOTE: Do NOT extend this list of arguments.
					offeredFeatures = []string{"xattrs", "delete", "compress"}
				}
			}

			// Find common features in both our type and offered type.
			commonFeatures := []string{}
			for _, ourFeature := range ourType.Features {
				if shared.StringInSlice(ourFeature, offeredFeatures) {
					commonFeatures = append(commonFeatures, ourFeature)
				}
			}

			// Append type with combined features.
			matchedTypes = append(matchedTypes, Type{
				FSType:   ourType.FSType,
				Features: commonFeatures,
			})
		}
	}

	if len(matchedTypes) < 1 {
		// No matching transport type found, generate an error with offered types and our types.
		offeredTypeStrings := make([]string, 0, len(offeredFSTypes))
		for _, offerFSType := range offeredFSTypes {
			offeredTypeStrings = append(offeredTypeStrings, offerFSType.String())
		}

		ourTypeStrings := make([]string, 0, len(ourTypes))
		for _, ourType := range ourTypes {
			ourTypeStrings = append(ourTypeStrings, ourType.FSType.String())
		}

		return matchedTypes, fmt.Errorf("No matching migration types found. Offered types: %v, our types: %v", offeredTypeStrings, ourTypeStrings)
	}

	return matchedTypes, nil
}

func progressWrapperRender(op *operations.Operation, key string, description string, progressInt int64, speedInt int64) {
	meta := op.Metadata()
	if meta == nil {
		meta = make(map[string]interface{})
	}

	progress := fmt.Sprintf("%s (%s/s)", units.GetByteSizeString(progressInt, 2), units.GetByteSizeString(speedInt, 2))
	if description != "" {
		progress = fmt.Sprintf("%s: %s (%s/s)", description, units.GetByteSizeString(progressInt, 2), units.GetByteSizeString(speedInt, 2))
	}

	if meta[key] != progress {
		meta[key] = progress
		op.UpdateMetadata(meta)
	}
}

// ProgressReader reports the read progress.
func ProgressReader(op *operations.Operation, key string, description string) func(io.ReadCloser) io.ReadCloser {
	return func(reader io.ReadCloser) io.ReadCloser {
		if op == nil {
			return reader
		}

		progress := func(progressInt int64, speedInt int64) {
			progressWrapperRender(op, key, description, progressInt, speedInt)
		}

		readPipe := &ioprogress.ProgressReader{
			ReadCloser: reader,
			Tracker: &ioprogress.ProgressTracker{
				Handler: progress,
			},
		}

		return readPipe
	}
}

// ProgressWriter reports the write progress.
func ProgressWriter(op *operations.Operation, key string, description string) func(io.WriteCloser) io.WriteCloser {
	return func(writer io.WriteCloser) io.WriteCloser {
		if op == nil {
			return writer
		}

		progress := func(progressInt int64, speedInt int64) {
			progressWrapperRender(op, key, description, progressInt, speedInt)
		}

		writePipe := &ioprogress.ProgressWriter{
			WriteCloser: writer,
			Tracker: &ioprogress.ProgressTracker{
				Handler: progress,
			},
		}

		return writePipe
	}
}

// ProgressTracker returns a migration I/O tracker
func ProgressTracker(op *operations.Operation, key string, description string) *ioprogress.ProgressTracker {
	progress := func(progressInt int64, speedInt int64) {
		progressWrapperRender(op, key, description, progressInt, speedInt)
	}

	tracker := &ioprogress.ProgressTracker{
		Handler: progress,
	}

	return tracker
}
<|endoftext|>


<commit_before>// Copyright (c) 2019-2020 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wallet

import (
	"context"
	"crypto/rand"
	"crypto/tls"
	"net"
	"time"

	"decred.org/cspp"
	"decred.org/cspp/coinjoin"
	"decred.org/dcrwallet/v2/errors"
	"decred.org/dcrwallet/v2/wallet/txauthor"
	"decred.org/dcrwallet/v2/wallet/txrules"
	"decred.org/dcrwallet/v2/wallet/txsizes"
	"decred.org/dcrwallet/v2/wallet/udb"
	"decred.org/dcrwallet/v2/wallet/walletdb"
	"github.com/decred/dcrd/dcrutil/v4"
	"github.com/decred/dcrd/wire"
	"github.com/decred/go-socks/socks"
	"golang.org/x/sync/errgroup"
)

// must be sorted large to small
var splitPoints = [...]dcrutil.Amount{
	1 << 36, // 687.19476736
	1 << 34, // 171.79869184
	1 << 32, // 042.94967296
	1 << 30, // 010.73741824
	1 << 28, // 002.68435456
	1 << 26, // 000.67108864
	1 << 24, // 000.16777216
	1 << 22, // 000.04194304
	1 << 20, // 000.01048576
	1 << 18, // 000.00262144
}

var splitSems = [len(splitPoints)]chan struct{}{}

func init() {
	for i := range splitSems {
		splitSems[i] = make(chan struct{}, 10)
	}
}

var (
	errNoSplitDenomination = errors.New("no suitable split denomination")
	errThrottledMixRequest = errors.New("throttled mix request for split denomination")
)

// DialFunc provides a method to dial a network connection.
// If the dialed network connection is secured by TLS, TLS
// configuration is provided by the method, not the caller.
type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)

func (w *Wallet) MixOutput(ctx context.Context, dialTLS DialFunc, csppserver string, output *wire.OutPoint, changeAccount, mixAccount, mixBranch uint32) error {
	op := errors.Opf("wallet.MixOutput(%v)", output)

	sdiff, err := w.NextStakeDifficulty(ctx)
	if err != nil {
		return errors.E(op, err)
	}

	defer w.holdUnlock().release()

	w.lockedOutpointMu.Lock()
	if _, exists := w.lockedOutpoints[outpoint{output.Hash, output.Index}]; exists {
		w.lockedOutpointMu.Unlock()
		err = errors.Errorf("output %v already locked", output)
		return errors.E(op, err)
	}

	var prevScript []byte
	var amount dcrutil.Amount
	err = walletdb.View(ctx, w.db, func(dbtx walletdb.ReadTx) error {
		txmgrNs := dbtx.ReadBucket(wtxmgrNamespaceKey)
		txDetails, err := w.txStore.TxDetails(txmgrNs, &output.Hash)
		if err != nil {
			return err
		}
		prevScript = txDetails.MsgTx.TxOut[output.Index].PkScript
		amount = dcrutil.Amount(txDetails.MsgTx.TxOut[output.Index].Value)
		return nil
	})
	if err != nil {
		w.lockedOutpointMu.Unlock()
		return errors.E(op, err)
	}
	w.lockedOutpoints[outpoint{output.Hash, output.Index}] = struct{}{}
	w.lockedOutpointMu.Unlock()

	defer func() {
		w.lockedOutpointMu.Lock()
		delete(w.lockedOutpoints, outpoint{output.Hash, output.Index})
		w.lockedOutpointMu.Unlock()
	}()

	var count int
	var mixValue, remValue dcrutil.Amount
	for i, v := range splitPoints {
		// When the sdiff is more than this mixed output amount, there
		// is a smaller common mixed amount with more pairing activity
		// (due to CoinShuffle++ participation from ticket buyers).
		// Skipping this amount and moving to the next smallest common
		// mixed amount will result in quicker pairings, or pairings
		// occurring at all.  The number of mixed outputs is capped to
		// prevent a single mix being overwhelmingly funded by a single
		// output, and to conserve memory resources.
		if i != len(splitPoints)-1 && v >= sdiff {
			continue
		}
		count = int(amount / v)
		if count > 4 {
			count = 4
		}
		if count > 0 {
			remValue = amount - dcrutil.Amount(count)*v
			mixValue = v
			select {
			case <-ctx.Done():
				return errors.E(op, ctx.Err())
			case splitSems[i] <- struct{}{}:
				defer func() { <-splitSems[i] }()
			default:
				return errThrottledMixRequest
			}
			break
		}
	}
	if mixValue == splitPoints[len(splitPoints)-1] {
		remValue = 0
	}
	if mixValue == 0 {
		err := errors.Errorf("output %v (%v): %w", output, amount, errNoSplitDenomination)
		return errors.E(op, err)
	}

	// Create change output from remaining value and contributed fee
	const P2PKHv0Len = 25
	feeRate := w.RelayFee()
	inScriptSizes := []int{txsizes.RedeemP2PKHSigScriptSize}
	outScriptSizes := make([]int, count)
	for i := range outScriptSizes {
		outScriptSizes[i] = P2PKHv0Len
	}
	size := txsizes.EstimateSerializeSizeFromScriptSizes(inScriptSizes, outScriptSizes, P2PKHv0Len)
	changeValue := remValue - txrules.FeeForSerializeSize(feeRate, size)
	var change *wire.TxOut
	var updates []func(walletdb.ReadWriteTx) error
	if !txrules.IsDustAmount(changeValue, P2PKHv0Len, feeRate) {
		persist := w.deferPersistReturnedChild(ctx, &updates)
		const accountName = "" // not used, so can be faked.
		addr, err := w.nextAddress(ctx, op, persist,
			accountName, changeAccount, udb.InternalBranch, WithGapPolicyIgnore())
		if err != nil {
			return errors.E(op, err)
		}
		changeScript, version, err := addressScript(addr)
		if err != nil {
			return errors.E(op, err)
		}
		change = &wire.TxOut{
			Value:    int64(changeValue),
			PkScript: changeScript,
			Version:  version,
		}
	}

	const (
		txVersion = 1
		locktime  = 0
		expiry    = 0
	)
	pairing := coinjoin.EncodeDesc(coinjoin.P2PKHv0, int64(mixValue), txVersion, locktime, expiry)
	ses, err := cspp.NewSession(rand.Reader, debugLog, pairing, count)
	if err != nil {
		return errors.E(op, err)
	}
	var conn net.Conn
	if dialTLS != nil {
		conn, err = dialTLS(ctx, "tcp", csppserver)
	} else {
		conn, err = tls.Dial("tcp", csppserver, nil)
	}
	if err != nil {
		return errors.E(op, err)
	}
	defer conn.Close()
	log.Infof("Dialed CSPPServer %v -> %v", conn.LocalAddr(), conn.RemoteAddr())

	log.Infof("Mixing output %v (%v)", output, amount)
	cj := w.newCsppJoin(ctx, change, mixValue, mixAccount, mixBranch, count)
	cj.addTxIn(prevScript, &wire.TxIn{
		PreviousOutPoint: *output,
		ValueIn:          int64(amount),
	})
	err = ses.DiceMix(ctx, conn, cj)
	if err != nil {
		return errors.E(op, err)
	}
	cjHash := cj.tx.TxHash()
	log.Infof("Completed CoinShuffle++ mix of output %v in transaction %v", output, &cjHash)

	var watch []wire.OutPoint
	w.lockedOutpointMu.Lock()
	err = walletdb.Update(ctx, w.db, func(dbtx walletdb.ReadWriteTx) error {
		for _, f := range updates {
			if err := f(dbtx); err != nil {
				return err
			}
		}
		rec, err := udb.NewTxRecordFromMsgTx(cj.tx, time.Now())
		if err != nil {
			return errors.E(op, err)
		}
		watch, err = w.processTransactionRecord(ctx, dbtx, rec, nil, nil)
		if err != nil {
			return err
		}
		return nil
	})
	w.lockedOutpointMu.Unlock()
	if err != nil {
		return errors.E(op, err)
	}
	n, _ := w.NetworkBackend()
	if n != nil {
		err = w.publishAndWatch(ctx, op, n, cj.tx, watch)
	}
	return err
}

// MixAccount individually mixes outputs of an account into standard
// denominations, creating newly mixed outputs for a mixed account.
//
// Due to performance concerns of timing out in a CoinShuffle++ run, this
// function may throttle how many of the outputs are mixed each call.
func (w *Wallet) MixAccount(ctx context.Context, dialTLS DialFunc, csppserver string, changeAccount, mixAccount, mixBranch uint32) error {
	const op errors.Op = "wallet.MixAccount"

	defer w.holdUnlock().release()

	_, tipHeight := w.MainChainTip(ctx)
	w.lockedOutpointMu.Lock()
	var credits []Input
	err := walletdb.View(ctx, w.db, func(dbtx walletdb.ReadTx) error {
		var err error
		credits, err = w.findEligibleOutputs(dbtx, changeAccount, 1, tipHeight)
		return err
	})
	if err != nil {
		w.lockedOutpointMu.Unlock()
		return errors.E(op, err)
	}
	validCredits := credits[:0]
	for i := range credits {
		amount := dcrutil.Amount(credits[i].PrevOut.Value)
		if amount <= splitPoints[len(splitPoints)-1] {
			continue
		}
		validCredits = append(validCredits, credits[i])
	}
	credits = validCredits
	shuffle(len(credits), func(i, j int) {
		credits[i], credits[j] = credits[j], credits[i]
	})
	if len(credits) > 32 { // simple throttle
		credits = credits[:32]
	}
	w.lockedOutpointMu.Unlock()

	var g errgroup.Group
	for i := range credits {
		op := &credits[i].OutPoint
		g.Go(func() error {
			err := w.MixOutput(ctx, dialTLS, csppserver, op, changeAccount, mixAccount, mixBranch)
			if errors.Is(err, errThrottledMixRequest) {
				return nil
			}
			if errors.Is(err, errNoSplitDenomination) {
				return nil
			}
			if errors.Is(err, socks.ErrPoolMaxConnections) {
				return nil
			}
			return err
		})
	}
	err = g.Wait()
	if err != nil {
		return errors.E(op, err)
	}
	return nil
}

// randomInputSource wraps an InputSource to randomly pick UTXOs.
// This involves reading all UTXOs from the underlying source into memory.
func randomInputSource(source txauthor.InputSource) txauthor.InputSource {
	all, err := source(dcrutil.MaxAmount)
	if err == nil {
		shuffleUTXOs(all)
	}
	var n int
	var tot dcrutil.Amount
	return func(target dcrutil.Amount) (*txauthor.InputDetail, error) {
		if err != nil {
			return nil, err
		}
		if all.Amount <= target {
			return all, nil
		}
		for n < len(all.Inputs) {
			tot += dcrutil.Amount(all.Inputs[n].ValueIn)
			n++
			if tot >= target {
				break
			}
		}
		selected := &txauthor.InputDetail{
			Amount:            tot,
			Inputs:            all.Inputs[:n],
			Scripts:           all.Scripts[:n],
			RedeemScriptSizes: all.RedeemScriptSizes[:n],
		}
		return selected, nil
	}
}

// PossibleCoinJoin tests if a transaction may be a CSPP-mixed transaction.
// It can return false positives, as one can create a tx which looks like a
// coinjoin tx, although it isn't.
func PossibleCoinJoin(tx *wire.MsgTx) (isMix bool, mixDenom int64, mixCount uint32) {
	if len(tx.TxOut) < 3 || len(tx.TxIn) < 3 {
		return false, 0, 0
	}

	numberOfOutputs := len(tx.TxOut)
	numberOfInputs := len(tx.TxIn)

	mixedOuts := make(map[int64]uint32)
	scripts := make(map[string]int)
	for _, o := range tx.TxOut {
		scripts[string(o.PkScript)]++
		if scripts[string(o.PkScript)] > 1 {
			return false, 0, 0
		}
		val := o.Value
		// Multiple zero valued outputs do not count as a coinjoin mix.
		if val == 0 {
			continue
		}
		mixedOuts[val]++
	}

	for val, count := range mixedOuts {
		if count < 3 {
			continue
		}
		if val > mixDenom {
			mixDenom = val
			mixCount = count
		}

		outputsWithNotSameAmount := uint32(numberOfOutputs) - count
		if outputsWithNotSameAmount > uint32(numberOfInputs) {
			return false, 0, 0
		}
	}

	isMix = mixCount >= uint32(len(tx.TxOut)/2)
	return
}
<commit_msg>Consider tx fees when choosing a mix amount and count<commit_after>// Copyright (c) 2019-2020 The Decred developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wallet

import (
	"context"
	"crypto/rand"
	"crypto/tls"
	"net"
	"time"

	"decred.org/cspp"
	"decred.org/cspp/coinjoin"
	"decred.org/dcrwallet/v2/errors"
	"decred.org/dcrwallet/v2/wallet/txauthor"
	"decred.org/dcrwallet/v2/wallet/txrules"
	"decred.org/dcrwallet/v2/wallet/txsizes"
	"decred.org/dcrwallet/v2/wallet/udb"
	"decred.org/dcrwallet/v2/wallet/walletdb"
	"github.com/decred/dcrd/dcrutil/v4"
	"github.com/decred/dcrd/wire"
	"github.com/decred/go-socks/socks"
	"golang.org/x/sync/errgroup"
)

// must be sorted large to small
var splitPoints = [...]dcrutil.Amount{
	1 << 36, // 687.19476736
	1 << 34, // 171.79869184
	1 << 32, // 042.94967296
	1 << 30, // 010.73741824
	1 << 28, // 002.68435456
	1 << 26, // 000.67108864
	1 << 24, // 000.16777216
	1 << 22, // 000.04194304
	1 << 20, // 000.01048576
	1 << 18, // 000.00262144
}

var splitSems = [len(splitPoints)]chan struct{}{}

func init() {
	for i := range splitSems {
		splitSems[i] = make(chan struct{}, 10)
	}
}

var (
	errNoSplitDenomination = errors.New("no suitable split denomination")
	errThrottledMixRequest = errors.New("throttled mix request for split denomination")
)

// DialFunc provides a method to dial a network connection.
// If the dialed network connection is secured by TLS, TLS
// configuration is provided by the method, not the caller.
type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)

func (w *Wallet) MixOutput(ctx context.Context, dialTLS DialFunc, csppserver string, output *wire.OutPoint, changeAccount, mixAccount, mixBranch uint32) error {
	op := errors.Opf("wallet.MixOutput(%v)", output)

	sdiff, err := w.NextStakeDifficulty(ctx)
	if err != nil {
		return errors.E(op, err)
	}

	defer w.holdUnlock().release()

	w.lockedOutpointMu.Lock()
	if _, exists := w.lockedOutpoints[outpoint{output.Hash, output.Index}]; exists {
		w.lockedOutpointMu.Unlock()
		err = errors.Errorf("output %v already locked", output)
		return errors.E(op, err)
	}

	var prevScript []byte
	var amount dcrutil.Amount
	err = walletdb.View(ctx, w.db, func(dbtx walletdb.ReadTx) error {
		txmgrNs := dbtx.ReadBucket(wtxmgrNamespaceKey)
		txDetails, err := w.txStore.TxDetails(txmgrNs, &output.Hash)
		if err != nil {
			return err
		}
		prevScript = txDetails.MsgTx.TxOut[output.Index].PkScript
		amount = dcrutil.Amount(txDetails.MsgTx.TxOut[output.Index].Value)
		return nil
	})
	if err != nil {
		w.lockedOutpointMu.Unlock()
		return errors.E(op, err)
	}
	w.lockedOutpoints[outpoint{output.Hash, output.Index}] = struct{}{}
	w.lockedOutpointMu.Unlock()

	defer func() {
		w.lockedOutpointMu.Lock()
		delete(w.lockedOutpoints, outpoint{output.Hash, output.Index})
		w.lockedOutpointMu.Unlock()
	}()

	var i, count int
	var mixValue, remValue, changeValue dcrutil.Amount
	var feeRate = w.RelayFee()
SplitPoints:
	for i = 0; i < len(splitPoints); i++ {
		last := i == len(splitPoints)-1

		// When the sdiff is more than this mixed output amount, there
		// is a smaller common mixed amount with more pairing activity
		// (due to CoinShuffle++ participation from ticket buyers).
		// Skipping this amount and moving to the next smallest common
		// mixed amount will result in quicker pairings, or pairings
		// occurring at all.  The number of mixed outputs is capped to
		// prevent a single mix being overwhelmingly funded by a single
		// output, and to conserve memory resources.
		if !last && mixValue >= sdiff {
			continue
		}

		mixValue = splitPoints[i]
		count = int(amount / mixValue)
		if count > 4 {
			count = 4
		}
		for ; count > 0; count-- {
			remValue = amount - dcrutil.Amount(count)*mixValue
			if remValue < 0 {
				continue
			}

			// Determine required fee and change value, if possible.
			// No change is ever included when mixing at the
			// smallest amount.
			const P2PKHv0Len = 25
			inScriptSizes := []int{txsizes.RedeemP2PKHSigScriptSize}
			outScriptSizes := make([]int, count)
			for i := range outScriptSizes {
				outScriptSizes[i] = P2PKHv0Len
			}
			size := txsizes.EstimateSerializeSizeFromScriptSizes(
				inScriptSizes, outScriptSizes, P2PKHv0Len)
			fee := txrules.FeeForSerializeSize(feeRate, size)
			changeValue = remValue - fee
			if last {
				changeValue = 0
			}
			if changeValue < 0 {
				// Determine required fee without a change
				// output.  A lower mix count or amount is
				// required if the fee is still not payable.
				size = txsizes.EstimateSerializeSizeFromScriptSizes(
					inScriptSizes, outScriptSizes, 0)
				fee = txrules.FeeForSerializeSize(feeRate, size)
				if remValue < fee {
					continue
				}
				changeValue = remValue - fee
			}
			if txrules.IsDustAmount(changeValue, P2PKHv0Len, feeRate) {
				changeValue = 0
			}

			break SplitPoints
		}
	}
	if i == len(splitPoints) {
		err := errors.Errorf("output %v (%v): %w", output, amount, errNoSplitDenomination)
		return errors.E(op, err)
	}
	select {
	case <-ctx.Done():
		return errors.E(op, ctx.Err())
	case splitSems[i] <- struct{}{}:
		defer func() { <-splitSems[i] }()
	default:
		return errThrottledMixRequest
	}

	var change *wire.TxOut
	var updates []func(walletdb.ReadWriteTx) error
	if changeValue > 0 {
		persist := w.deferPersistReturnedChild(ctx, &updates)
		const accountName = "" // not used, so can be faked.
		addr, err := w.nextAddress(ctx, op, persist,
			accountName, changeAccount, udb.InternalBranch, WithGapPolicyIgnore())
		if err != nil {
			return errors.E(op, err)
		}
		changeScript, version, err := addressScript(addr)
		if err != nil {
			return errors.E(op, err)
		}
		change = &wire.TxOut{
			Value:    int64(changeValue),
			PkScript: changeScript,
			Version:  version,
		}
	}

	const (
		txVersion = 1
		locktime  = 0
		expiry    = 0
	)
	pairing := coinjoin.EncodeDesc(coinjoin.P2PKHv0, int64(mixValue), txVersion, locktime, expiry)
	ses, err := cspp.NewSession(rand.Reader, debugLog, pairing, count)
	if err != nil {
		return errors.E(op, err)
	}
	var conn net.Conn
	if dialTLS != nil {
		conn, err = dialTLS(ctx, "tcp", csppserver)
	} else {
		conn, err = tls.Dial("tcp", csppserver, nil)
	}
	if err != nil {
		return errors.E(op, err)
	}
	defer conn.Close()
	log.Infof("Dialed CSPPServer %v -> %v", conn.LocalAddr(), conn.RemoteAddr())

	log.Infof("Mixing output %v (%v)", output, amount)
	cj := w.newCsppJoin(ctx, change, mixValue, mixAccount, mixBranch, count)
	cj.addTxIn(prevScript, &wire.TxIn{
		PreviousOutPoint: *output,
		ValueIn:          int64(amount),
	})
	err = ses.DiceMix(ctx, conn, cj)
	if err != nil {
		return errors.E(op, err)
	}
	cjHash := cj.tx.TxHash()
	log.Infof("Completed CoinShuffle++ mix of output %v in transaction %v", output, &cjHash)

	var watch []wire.OutPoint
	w.lockedOutpointMu.Lock()
	err = walletdb.Update(ctx, w.db, func(dbtx walletdb.ReadWriteTx) error {
		for _, f := range updates {
			if err := f(dbtx); err != nil {
				return err
			}
		}
		rec, err := udb.NewTxRecordFromMsgTx(cj.tx, time.Now())
		if err != nil {
			return errors.E(op, err)
		}
		watch, err = w.processTransactionRecord(ctx, dbtx, rec, nil, nil)
		if err != nil {
			return err
		}
		return nil
	})
	w.lockedOutpointMu.Unlock()
	if err != nil {
		return errors.E(op, err)
	}
	n, _ := w.NetworkBackend()
	if n != nil {
		err = w.publishAndWatch(ctx, op, n, cj.tx, watch)
	}
	return err
}

// MixAccount individually mixes outputs of an account into standard
// denominations, creating newly mixed outputs for a mixed account.
//
// Due to performance concerns of timing out in a CoinShuffle++ run, this
// function may throttle how many of the outputs are mixed each call.
func (w *Wallet) MixAccount(ctx context.Context, dialTLS DialFunc, csppserver string, changeAccount, mixAccount, mixBranch uint32) error {
	const op errors.Op = "wallet.MixAccount"

	defer w.holdUnlock().release()

	_, tipHeight := w.MainChainTip(ctx)
	w.lockedOutpointMu.Lock()
	var credits []Input
	err := walletdb.View(ctx, w.db, func(dbtx walletdb.ReadTx) error {
		var err error
		credits, err = w.findEligibleOutputs(dbtx, changeAccount, 1, tipHeight)
		return err
	})
	if err != nil {
		w.lockedOutpointMu.Unlock()
		return errors.E(op, err)
	}
	validCredits := credits[:0]
	for i := range credits {
		amount := dcrutil.Amount(credits[i].PrevOut.Value)
		if amount <= splitPoints[len(splitPoints)-1] {
			continue
		}
		validCredits = append(validCredits, credits[i])
	}
	credits = validCredits
	shuffle(len(credits), func(i, j int) {
		credits[i], credits[j] = credits[j], credits[i]
	})
	if len(credits) > 32 { // simple throttle
		credits = credits[:32]
	}
	w.lockedOutpointMu.Unlock()

	var g errgroup.Group
	for i := range credits {
		op := &credits[i].OutPoint
		g.Go(func() error {
			err := w.MixOutput(ctx, dialTLS, csppserver, op, changeAccount, mixAccount, mixBranch)
			if errors.Is(err, errThrottledMixRequest) {
				return nil
			}
			if errors.Is(err, errNoSplitDenomination) {
				return nil
			}
			if errors.Is(err, socks.ErrPoolMaxConnections) {
				return nil
			}
			return err
		})
	}
	err = g.Wait()
	if err != nil {
		return errors.E(op, err)
	}
	return nil
}

// randomInputSource wraps an InputSource to randomly pick UTXOs.
// This involves reading all UTXOs from the underlying source into memory.
func randomInputSource(source txauthor.InputSource) txauthor.InputSource {
	all, err := source(dcrutil.MaxAmount)
	if err == nil {
		shuffleUTXOs(all)
	}
	var n int
	var tot dcrutil.Amount
	return func(target dcrutil.Amount) (*txauthor.InputDetail, error) {
		if err != nil {
			return nil, err
		}
		if all.Amount <= target {
			return all, nil
		}
		for n < len(all.Inputs) {
			tot += dcrutil.Amount(all.Inputs[n].ValueIn)
			n++
			if tot >= target {
				break
			}
		}
		selected := &txauthor.InputDetail{
			Amount:            tot,
			Inputs:            all.Inputs[:n],
			Scripts:           all.Scripts[:n],
			RedeemScriptSizes: all.RedeemScriptSizes[:n],
		}
		return selected, nil
	}
}

// PossibleCoinJoin tests if a transaction may be a CSPP-mixed transaction.
// It can return false positives, as one can create a tx which looks like a
// coinjoin tx, although it isn't.
func PossibleCoinJoin(tx *wire.MsgTx) (isMix bool, mixDenom int64, mixCount uint32) {
	if len(tx.TxOut) < 3 || len(tx.TxIn) < 3 {
		return false, 0, 0
	}

	numberOfOutputs := len(tx.TxOut)
	numberOfInputs := len(tx.TxIn)

	mixedOuts := make(map[int64]uint32)
	scripts := make(map[string]int)
	for _, o := range tx.TxOut {
		scripts[string(o.PkScript)]++
		if scripts[string(o.PkScript)] > 1 {
			return false, 0, 0
		}
		val := o.Value
		// Multiple zero valued outputs do not count as a coinjoin mix.
		if val == 0 {
			continue
		}
		mixedOuts[val]++
	}

	for val, count := range mixedOuts {
		if count < 3 {
			continue
		}
		if val > mixDenom {
			mixDenom = val
			mixCount = count
		}

		outputsWithNotSameAmount := uint32(numberOfOutputs) - count
		if outputsWithNotSameAmount > uint32(numberOfInputs) {
			return false, 0, 0
		}
	}

	isMix = mixCount >= uint32(len(tx.TxOut)/2)
	return
}
<|endoftext|>


<commit_before>package library

import (
	"encoding/json"
	"github.com/bitly/go-nsq"
	"github.com/nytlabs/streamtools/st/blocks"
	"github.com/nytlabs/streamtools/st/util"
)

// specify those channels we're going to use to communicate with streamtools
type FromNSQ struct {
	blocks.Block
	queryrule   chan chan interface{}
	inrule      chan interface{}
	out         chan interface{}
	quit        chan interface{}
	topic       string
	channel     string
	lookupdAddr string
	maxInFlight int
}

// a bit of boilerplate for streamtools
func NewFromNSQ() blocks.BlockInterface {
	return &FromNSQ{}
}

func (b *FromNSQ) Setup() {
	b.Kind = "FromNSQ"
	b.inrule = b.InRoute("rule")
	b.queryrule = b.QueryRoute("rule")
	b.quit = b.Quit()
	b.out = b.Broadcast()
}

type readWriteHandler struct {
	toOut   chan interface{}
	toError chan error
}

func (self readWriteHandler) HandleMessage(message *nsq.Message) error {
	var msg interface{}
	err := json.Unmarshal(message.Body, &msg)
	if err != nil {
		self.toError <- err
		return err
	}
	self.toOut <- msg
	return nil
}

// connects to an NSQ topic and emits each message into streamtools.
func (b *FromNSQ) Run() {
	var reader *nsq.Reader
	toOut := make(chan interface{})
	toError := make(chan error)

	for {
		select {
		case msg := <-toOut:
			b.out <- &msg
		case err := <-toError:
			b.Error(err)
		case ruleI := <-b.inrule:
			// convert message to a map of string interfaces
			// aka keys are strings, values are empty interfaces
			rule := ruleI.(map[string]interface{})

			topic, err := util.ParseString(rule, "ReadTopic")
			if err != nil {
				b.Error(err)
			}

			lookupdAddr, err := util.ParseString(rule, "LookupdAddr")
			if err != nil {
				b.Error(err)
			}
			maxInFlight, err := util.ParseInt(rule, "MaxInFlight")
			if err != nil {
				b.Error(err)
			}
			channel, err := util.ParseString(rule, "ReadChannel")
			if err != nil {
				b.Error(err)
			}

			reader, err := nsq.NewReader(topic, channel)
			if err != nil {
				b.Error(err)
			}
			reader.SetMaxInFlight(maxInFlight)

			h := readWriteHandler{toOut, toError}
			reader.AddHandler(h)

			err = reader.ConnectToLookupd(lookupdAddr)
			if err != nil {
				b.Error(err)
			}

			b.topic = topic
			b.channel = channel
			b.maxInFlight = maxInFlight
			b.lookupdAddr = lookupdAddr

		case <-b.quit:
			if reader != nil {
				reader.Stop()
			}
			return
		case c := <-b.queryrule:
			c <- map[string]interface{}{
				"ReadTopic":   b.topic,
				"ReadChannel": b.channel,
				"LookupdAddr": b.lookupdAddr,
				"MaxInFlight": b.maxInFlight,
			}
		}
	}
}
<commit_msg>maxInFlight wants to be an int but is a float64 because json<commit_after>package library

import (
	"encoding/json"
	"github.com/bitly/go-nsq"
	"github.com/nytlabs/streamtools/st/blocks"
	"github.com/nytlabs/streamtools/st/util"
)

// specify those channels we're going to use to communicate with streamtools
type FromNSQ struct {
	blocks.Block
	queryrule   chan chan interface{}
	inrule      chan interface{}
	out         chan interface{}
	quit        chan interface{}
	topic       string
	channel     string
	lookupdAddr string
	maxInFlight float64
}

// a bit of boilerplate for streamtools
func NewFromNSQ() blocks.BlockInterface {
	return &FromNSQ{}
}

func (b *FromNSQ) Setup() {
	b.Kind = "FromNSQ"
	b.inrule = b.InRoute("rule")
	b.queryrule = b.QueryRoute("rule")
	b.quit = b.Quit()
	b.out = b.Broadcast()
}

type readWriteHandler struct {
	toOut   chan interface{}
	toError chan error
}

func (self readWriteHandler) HandleMessage(message *nsq.Message) error {
	var msg interface{}
	err := json.Unmarshal(message.Body, &msg)
	if err != nil {
		self.toError <- err
		return err
	}
	self.toOut <- msg
	return nil
}

// connects to an NSQ topic and emits each message into streamtools.
func (b *FromNSQ) Run() {
	var reader *nsq.Reader
	toOut := make(chan interface{})
	toError := make(chan error)

	for {
		select {
		case msg := <-toOut:
			b.out <- &msg
		case err := <-toError:
			b.Error(err)
		case ruleI := <-b.inrule:
			// convert message to a map of string interfaces
			// aka keys are strings, values are empty interfaces
			rule := ruleI.(map[string]interface{})

			topic, err := util.ParseString(rule, "ReadTopic")
			if err != nil {
				b.Error(err)
			}

			lookupdAddr, err := util.ParseString(rule, "LookupdAddr")
			if err != nil {
				b.Error(err)
			}
			maxInFlight, err := util.ParseFloat(rule, "MaxInFlight")
			if err != nil {
				b.Error(err)
			}
			channel, err := util.ParseString(rule, "ReadChannel")
			if err != nil {
				b.Error(err)
			}

			reader, err := nsq.NewReader(topic, channel)
			if err != nil {
				b.Error(err)
			}
			reader.SetMaxInFlight(int(maxInFlight))

			h := readWriteHandler{toOut, toError}
			reader.AddHandler(h)

			err = reader.ConnectToLookupd(lookupdAddr)
			if err != nil {
				b.Error(err)
			}

			b.topic = topic
			b.channel = channel
			b.maxInFlight = maxInFlight
			b.lookupdAddr = lookupdAddr

		case <-b.quit:
			if reader != nil {
				reader.Stop()
			}
			return
		case c := <-b.queryrule:
			c <- map[string]interface{}{
				"ReadTopic":   b.topic,
				"ReadChannel": b.channel,
				"LookupdAddr": b.lookupdAddr,
				"MaxInFlight": b.maxInFlight,
			}
		}
	}
}
<|endoftext|>


<commit_before>package library

import (
	"encoding/json"
	"github.com/bitly/go-nsq"
	"github.com/nytlabs/streamtools/st/blocks"
	"github.com/nytlabs/streamtools/st/util"
)

// specify those channels we're going to use to communicate with streamtools
type FromNSQ struct {
	blocks.Block
	queryrule   chan chan interface{}
	inrule      chan interface{}
	out         chan interface{}
	quit        chan interface{}
	topic       string
	channel     string
	lookupdAddr string
	maxInFlight float64
}

// a bit of boilerplate for streamtools
func NewFromNSQ() blocks.BlockInterface {
	return &FromNSQ{}
}

func (b *FromNSQ) Setup() {
	b.Kind = "FromNSQ"
	b.inrule = b.InRoute("rule")
	b.queryrule = b.QueryRoute("rule")
	b.quit = b.Quit()
	b.out = b.Broadcast()
}

type readWriteHandler struct {
	toOut   chan interface{}
	toError chan error
}

func (self readWriteHandler) HandleMessage(message *nsq.Message) error {
	var msg interface{}
	err := json.Unmarshal(message.Body, &msg)
	if err != nil {
		self.toError <- err
		return err
	}
	self.toOut <- msg
	return nil
}

// connects to an NSQ topic and emits each message into streamtools.
func (b *FromNSQ) Run() {
	var reader *nsq.Reader
	toOut := make(chan interface{})
	toError := make(chan error)

	for {
		select {
		case msg := <-toOut:
			b.out <- &msg
		case err := <-toError:
			b.Error(err)
		case ruleI := <-b.inrule:
			// convert message to a map of string interfaces
			// aka keys are strings, values are empty interfaces
			rule := ruleI.(map[string]interface{})

			topic, err := util.ParseString(rule, "ReadTopic")
			if err != nil {
				b.Error(err)
			}

			lookupdAddr, err := util.ParseString(rule, "LookupdAddr")
			if err != nil {
				b.Error(err)
			}
			maxInFlight, err := util.ParseFloat(rule, "MaxInFlight")
			if err != nil {
				b.Error(err)
			}
			channel, err := util.ParseString(rule, "ReadChannel")
			if err != nil {
				b.Error(err)
			}

			reader, err := nsq.NewReader(topic, channel)
			if err != nil {
				b.Error(err)
			}
			reader.SetMaxInFlight(int(maxInFlight))

			h := readWriteHandler{toOut, toError}
			reader.AddHandler(h)

			err = reader.ConnectToLookupd(lookupdAddr)
			if err != nil {
				b.Error(err)
			}

			b.topic = topic
			b.channel = channel
			b.maxInFlight = maxInFlight
			b.lookupdAddr = lookupdAddr

		case <-b.quit:
			if reader != nil {
				reader.Stop()
			}
			return
		case c := <-b.queryrule:
			c <- map[string]interface{}{
				"ReadTopic":   b.topic,
				"ReadChannel": b.channel,
				"LookupdAddr": b.lookupdAddr,
				"MaxInFlight": b.maxInFlight,
			}
		}
	}
}
<commit_msg>fix for NSQ<commit_after>package library

import (
	"encoding/json"
	"github.com/bitly/go-nsq"
	"github.com/nytlabs/streamtools/st/blocks"
	"github.com/nytlabs/streamtools/st/util"
)

// specify those channels we're going to use to communicate with streamtools
type FromNSQ struct {
	blocks.Block
	queryrule   chan chan interface{}
	inrule      chan interface{}
	out         chan interface{}
	quit        chan interface{}
	topic       string
	channel     string
	lookupdAddr string
	maxInFlight float64
}

// a bit of boilerplate for streamtools
func NewFromNSQ() blocks.BlockInterface {
	return &FromNSQ{}
}

func (b *FromNSQ) Setup() {
	b.Kind = "FromNSQ"
	b.inrule = b.InRoute("rule")
	b.queryrule = b.QueryRoute("rule")
	b.quit = b.Quit()
	b.out = b.Broadcast()
}

type readWriteHandler struct {
	toOut   chan interface{}
	toError chan error
}

func (self readWriteHandler) HandleMessage(message *nsq.Message) error {
	var msg interface{}
	err := json.Unmarshal(message.Body, &msg)
	if err != nil {
		self.toError <- err
		return err
	}
	self.toOut <- msg
	return nil
}

// connects to an NSQ topic and emits each message into streamtools.
func (b *FromNSQ) Run() {
	var reader *nsq.Reader
	toOut := make(chan interface{})
	toError := make(chan error)

	for {
		select {
		case msg := <-toOut:
			b.out <- msg
		case err := <-toError:
			b.Error(err)
		case ruleI := <-b.inrule:
			// convert message to a map of string interfaces
			// aka keys are strings, values are empty interfaces
			rule := ruleI.(map[string]interface{})

			topic, err := util.ParseString(rule, "ReadTopic")
			if err != nil {
				b.Error(err)
			}

			lookupdAddr, err := util.ParseString(rule, "LookupdAddr")
			if err != nil {
				b.Error(err)
			}
			maxInFlight, err := util.ParseFloat(rule, "MaxInFlight")
			if err != nil {
				b.Error(err)
			}
			channel, err := util.ParseString(rule, "ReadChannel")
			if err != nil {
				b.Error(err)
			}

			reader, err := nsq.NewReader(topic, channel)
			if err != nil {
				b.Error(err)
			}
			reader.SetMaxInFlight(int(maxInFlight))

			h := readWriteHandler{toOut, toError}
			reader.AddHandler(h)

			err = reader.ConnectToLookupd(lookupdAddr)
			if err != nil {
				b.Error(err)
			}

			b.topic = topic
			b.channel = channel
			b.maxInFlight = maxInFlight
			b.lookupdAddr = lookupdAddr

		case <-b.quit:
			if reader != nil {
				reader.Stop()
			}
			return
		case c := <-b.queryrule:
			c <- map[string]interface{}{
				"ReadTopic":   b.topic,
				"ReadChannel": b.channel,
				"LookupdAddr": b.lookupdAddr,
				"MaxInFlight": b.maxInFlight,
			}
		}
	}
}
<|endoftext|>


<commit_before>package topology

import (
	_ "fmt"
	"pkg/storage"
	"strconv"
)

type DataNode struct {
	NodeImpl
	volumes   map[storage.VolumeId]storage.VolumeInfo
	Ip        string
	Port      int
	PublicUrl string
	LastSeen  int64 // unix time in seconds
	Dead    bool
}

func NewDataNode(id string) *DataNode {
	s := &DataNode{}
	s.id = NodeId(id)
	s.nodeType = "DataNode"
	s.volumes = make(map[storage.VolumeId]storage.VolumeInfo)
  s.NodeImpl.value = s
	return s
}
func (dn *DataNode) AddOrUpdateVolume(v storage.VolumeInfo) {
	if _, ok := dn.volumes[v.Id]; !ok {
		dn.volumes[v.Id] = v
		dn.UpAdjustActiveVolumeCountDelta(1)
		dn.UpAdjustMaxVolumeId(v.Id)
	} else {
		dn.volumes[v.Id] = v
	}
}
func (dn *DataNode) GetTopology() *Topology {
	p := dn.parent
	for p.Parent() != nil {
		p = p.Parent()
	}
	t := p.(*Topology)
	return t
}
func (dn *DataNode) MatchLocation(ip string, port int) bool {
	return dn.Ip == ip && dn.Port == port
}
func (dn *DataNode) Url() string {
  return dn.Ip + strconv.Itoa(dn.Port)
}

func (dn *DataNode) ToMap() interface{} {
	ret := make(map[string]interface{})
	ret["Url"] = dn.Url()
	ret["Volumes"] = dn.GetActiveVolumeCount()
	ret["Max"] = dn.GetMaxVolumeCount()
	ret["Free"] = dn.FreeSpace()
	ret["PublicUrl"] = dn.PublicUrl
	return ret
}
<commit_msg>url error<commit_after>package topology

import (
	_ "fmt"
	"pkg/storage"
	"strconv"
)

type DataNode struct {
	NodeImpl
	volumes   map[storage.VolumeId]storage.VolumeInfo
	Ip        string
	Port      int
	PublicUrl string
	LastSeen  int64 // unix time in seconds
	Dead    bool
}

func NewDataNode(id string) *DataNode {
	s := &DataNode{}
	s.id = NodeId(id)
	s.nodeType = "DataNode"
	s.volumes = make(map[storage.VolumeId]storage.VolumeInfo)
  s.NodeImpl.value = s
	return s
}
func (dn *DataNode) AddOrUpdateVolume(v storage.VolumeInfo) {
	if _, ok := dn.volumes[v.Id]; !ok {
		dn.volumes[v.Id] = v
		dn.UpAdjustActiveVolumeCountDelta(1)
		dn.UpAdjustMaxVolumeId(v.Id)
	} else {
		dn.volumes[v.Id] = v
	}
}
func (dn *DataNode) GetTopology() *Topology {
	p := dn.parent
	for p.Parent() != nil {
		p = p.Parent()
	}
	t := p.(*Topology)
	return t
}
func (dn *DataNode) MatchLocation(ip string, port int) bool {
	return dn.Ip == ip && dn.Port == port
}
func (dn *DataNode) Url() string {
  return dn.Ip + ":" + strconv.Itoa(dn.Port)
}

func (dn *DataNode) ToMap() interface{} {
	ret := make(map[string]interface{})
	ret["Url"] = dn.Url()
	ret["Volumes"] = dn.GetActiveVolumeCount()
	ret["Max"] = dn.GetMaxVolumeCount()
	ret["Free"] = dn.FreeSpace()
	ret["PublicUrl"] = dn.PublicUrl
	return ret
}
<|endoftext|>


<commit_before>package remote_storage

import (
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/remote_pb"
	"github.com/golang/protobuf/proto"
	"io"
	"strings"
	"sync"
)

func ParseLocationName(remote string) (locationName string) {
	if strings.HasSuffix(string(remote), "/") {
		remote = remote[:len(remote)-1]
	}
	parts := strings.SplitN(string(remote), "/", 2)
	if len(parts) >= 1 {
		return parts[0]
	}
	return
}

func parseBucketLocation(remote string) (loc *remote_pb.RemoteStorageLocation) {
	loc = &remote_pb.RemoteStorageLocation{}
	if strings.HasSuffix(string(remote), "/") {
		remote = remote[:len(remote)-1]
	}
	parts := strings.SplitN(string(remote), "/", 3)
	if len(parts) >= 1 {
		loc.Name = parts[0]
	}
	if len(parts) >= 2 {
		loc.Bucket = parts[1]
	}
	loc.Path = string(remote[len(loc.Name)+1+len(loc.Bucket):])
	if loc.Path == "" {
		loc.Path = "/"
	}
	return
}

func parseNoBucketLocation(remote string) (loc *remote_pb.RemoteStorageLocation) {
	loc = &remote_pb.RemoteStorageLocation{}
	if strings.HasSuffix(string(remote), "/") {
		remote = remote[:len(remote)-1]
	}
	parts := strings.SplitN(string(remote), "/", 2)
	if len(parts) >= 1 {
		loc.Name = parts[0]
	}
	loc.Path = string(remote[len(loc.Name):])
	if loc.Path == "" {
		loc.Path = "/"
	}
	return
}

func FormatLocation(loc *remote_pb.RemoteStorageLocation) string {
	return fmt.Sprintf("%s/%s%s", loc.Name, loc.Bucket, loc.Path)
}

type VisitFunc func(dir string, name string, isDirectory bool, remoteEntry *filer_pb.RemoteEntry) error

type RemoteStorageClient interface {
	Traverse(loc *remote_pb.RemoteStorageLocation, visitFn VisitFunc) error
	ReadFile(loc *remote_pb.RemoteStorageLocation, offset int64, size int64) (data []byte, err error)
	WriteDirectory(loc *remote_pb.RemoteStorageLocation, entry *filer_pb.Entry) (err error)
	RemoveDirectory(loc *remote_pb.RemoteStorageLocation) (err error)
	WriteFile(loc *remote_pb.RemoteStorageLocation, entry *filer_pb.Entry, reader io.Reader) (remoteEntry *filer_pb.RemoteEntry, err error)
	UpdateFileMetadata(loc *remote_pb.RemoteStorageLocation, oldEntry *filer_pb.Entry, newEntry *filer_pb.Entry) (err error)
	DeleteFile(loc *remote_pb.RemoteStorageLocation) (err error)
}

type RemoteStorageClientMaker interface {
	Make(remoteConf *remote_pb.RemoteConf) (RemoteStorageClient, error)
	HasBucket() bool
}

type CachedRemoteStorageClient struct {
	*remote_pb.RemoteConf
	RemoteStorageClient
}

var (
	RemoteStorageClientMakers = make(map[string]RemoteStorageClientMaker)
	remoteStorageClients      = make(map[string]CachedRemoteStorageClient)
	remoteStorageClientsLock  sync.Mutex
)

func ParseRemoteLocation(remoteConfType string, remote string) (remoteStorageLocation *remote_pb.RemoteStorageLocation, err error) {
	maker, found := RemoteStorageClientMakers[remoteConfType]
	if !found {
		return nil, fmt.Errorf("remote storage type %s not found", remoteConfType)
	}

	if !maker.HasBucket() {
		return parseNoBucketLocation(remote), nil
	}
	return parseBucketLocation(remote), nil
}

func makeRemoteStorageClient(remoteConf *remote_pb.RemoteConf) (RemoteStorageClient, error) {
	maker, found := RemoteStorageClientMakers[remoteConf.Type]
	if !found {
		return nil, fmt.Errorf("remote storage type %s not found", remoteConf.Type)
	}
	return maker.Make(remoteConf)
}

func GetRemoteStorage(remoteConf *remote_pb.RemoteConf) (RemoteStorageClient, error) {
	remoteStorageClientsLock.Lock()
	defer remoteStorageClientsLock.Unlock()

	existingRemoteStorageClient, found := remoteStorageClients[remoteConf.Name]
	if found && proto.Equal(existingRemoteStorageClient.RemoteConf, remoteConf) {
		return existingRemoteStorageClient.RemoteStorageClient, nil
	}

	newRemoteStorageClient, err := makeRemoteStorageClient(remoteConf)
	if err != nil {
		return nil, fmt.Errorf("make remote storage client %s: %v", remoteConf.Name, err)
	}

	remoteStorageClients[remoteConf.Name] = CachedRemoteStorageClient{
		RemoteConf:          remoteConf,
		RemoteStorageClient: newRemoteStorageClient,
	}

	return newRemoteStorageClient, nil
}
<commit_msg>adjust formatting remote location<commit_after>package remote_storage

import (
	"fmt"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/pb/remote_pb"
	"github.com/golang/protobuf/proto"
	"io"
	"strings"
	"sync"
)

func ParseLocationName(remote string) (locationName string) {
	if strings.HasSuffix(string(remote), "/") {
		remote = remote[:len(remote)-1]
	}
	parts := strings.SplitN(string(remote), "/", 2)
	if len(parts) >= 1 {
		return parts[0]
	}
	return
}

func parseBucketLocation(remote string) (loc *remote_pb.RemoteStorageLocation) {
	loc = &remote_pb.RemoteStorageLocation{}
	if strings.HasSuffix(string(remote), "/") {
		remote = remote[:len(remote)-1]
	}
	parts := strings.SplitN(string(remote), "/", 3)
	if len(parts) >= 1 {
		loc.Name = parts[0]
	}
	if len(parts) >= 2 {
		loc.Bucket = parts[1]
	}
	loc.Path = string(remote[len(loc.Name)+1+len(loc.Bucket):])
	if loc.Path == "" {
		loc.Path = "/"
	}
	return
}

func parseNoBucketLocation(remote string) (loc *remote_pb.RemoteStorageLocation) {
	loc = &remote_pb.RemoteStorageLocation{}
	if strings.HasSuffix(string(remote), "/") {
		remote = remote[:len(remote)-1]
	}
	parts := strings.SplitN(string(remote), "/", 2)
	if len(parts) >= 1 {
		loc.Name = parts[0]
	}
	loc.Path = string(remote[len(loc.Name):])
	if loc.Path == "" {
		loc.Path = "/"
	}
	return
}

func FormatLocation(loc *remote_pb.RemoteStorageLocation) string {
	if loc.Bucket == "" {
		return fmt.Sprintf("%s%s", loc.Name, loc.Path)
	}
	return fmt.Sprintf("%s/%s%s", loc.Name, loc.Bucket, loc.Path)
}

type VisitFunc func(dir string, name string, isDirectory bool, remoteEntry *filer_pb.RemoteEntry) error

type RemoteStorageClient interface {
	Traverse(loc *remote_pb.RemoteStorageLocation, visitFn VisitFunc) error
	ReadFile(loc *remote_pb.RemoteStorageLocation, offset int64, size int64) (data []byte, err error)
	WriteDirectory(loc *remote_pb.RemoteStorageLocation, entry *filer_pb.Entry) (err error)
	RemoveDirectory(loc *remote_pb.RemoteStorageLocation) (err error)
	WriteFile(loc *remote_pb.RemoteStorageLocation, entry *filer_pb.Entry, reader io.Reader) (remoteEntry *filer_pb.RemoteEntry, err error)
	UpdateFileMetadata(loc *remote_pb.RemoteStorageLocation, oldEntry *filer_pb.Entry, newEntry *filer_pb.Entry) (err error)
	DeleteFile(loc *remote_pb.RemoteStorageLocation) (err error)
}

type RemoteStorageClientMaker interface {
	Make(remoteConf *remote_pb.RemoteConf) (RemoteStorageClient, error)
	HasBucket() bool
}

type CachedRemoteStorageClient struct {
	*remote_pb.RemoteConf
	RemoteStorageClient
}

var (
	RemoteStorageClientMakers = make(map[string]RemoteStorageClientMaker)
	remoteStorageClients      = make(map[string]CachedRemoteStorageClient)
	remoteStorageClientsLock  sync.Mutex
)

func ParseRemoteLocation(remoteConfType string, remote string) (remoteStorageLocation *remote_pb.RemoteStorageLocation, err error) {
	maker, found := RemoteStorageClientMakers[remoteConfType]
	if !found {
		return nil, fmt.Errorf("remote storage type %s not found", remoteConfType)
	}

	if !maker.HasBucket() {
		return parseNoBucketLocation(remote), nil
	}
	return parseBucketLocation(remote), nil
}

func makeRemoteStorageClient(remoteConf *remote_pb.RemoteConf) (RemoteStorageClient, error) {
	maker, found := RemoteStorageClientMakers[remoteConf.Type]
	if !found {
		return nil, fmt.Errorf("remote storage type %s not found", remoteConf.Type)
	}
	return maker.Make(remoteConf)
}

func GetRemoteStorage(remoteConf *remote_pb.RemoteConf) (RemoteStorageClient, error) {
	remoteStorageClientsLock.Lock()
	defer remoteStorageClientsLock.Unlock()

	existingRemoteStorageClient, found := remoteStorageClients[remoteConf.Name]
	if found && proto.Equal(existingRemoteStorageClient.RemoteConf, remoteConf) {
		return existingRemoteStorageClient.RemoteStorageClient, nil
	}

	newRemoteStorageClient, err := makeRemoteStorageClient(remoteConf)
	if err != nil {
		return nil, fmt.Errorf("make remote storage client %s: %v", remoteConf.Name, err)
	}

	remoteStorageClients[remoteConf.Name] = CachedRemoteStorageClient{
		RemoteConf:          remoteConf,
		RemoteStorageClient: newRemoteStorageClient,
	}

	return newRemoteStorageClient, nil
}
<|endoftext|>