<commit_before>package agent

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

	log "github.com/coreos/fleet/third_party/github.com/golang/glog"

	"github.com/coreos/fleet/event"
	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/registry"
	"github.com/coreos/fleet/sign"
	"github.com/coreos/fleet/systemd"
	"github.com/coreos/fleet/unit"
)

const (
	// TTL to use with all state pushed to Registry
	DefaultTTL = "30s"

	// Refresh TTLs at 1/2 the TTL length
	refreshInterval = 2
)

// The Agent owns all of the coordination between the Registry, the local
// Machine, and the local SystemdManager.
type Agent struct {
	registry      *registry.Registry
	events        *event.EventBus
	machine       *machine.Machine
	ttl           time.Duration
	systemdPrefix string
	// verifier is used to verify job payload. A nil one implies that
	// all payloads are accepted.
	verifier *sign.SignatureVerifier

	state   *AgentState
	systemd *systemd.SystemdManager

	// channel used to shutdown any open connections/channels the Agent holds
	stop chan bool
}

func New(registry *registry.Registry, events *event.EventBus, machine *machine.Machine, ttl, unitPrefix string, verifier *sign.SignatureVerifier) (*Agent, error) {
	ttldur, err := time.ParseDuration(ttl)
	if err != nil {
		return nil, err
	}

	state := NewState()
	mgr := systemd.NewSystemdManager(machine, unitPrefix)

	return &Agent{registry, events, machine, ttldur, unitPrefix, verifier, state, mgr, nil}, nil
}

// Access Agent's machine field
func (a *Agent) Machine() *machine.Machine {
	return a.machine
}

func (a *Agent) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct{ State *AgentState }{State: a.state})
}

// Trigger all async processes the Agent intends to run
func (a *Agent) Run() {
	a.stop = make(chan bool)

	handler := NewEventHandler(a)
	a.events.AddListener("agent", a.machine, handler)

	go a.systemd.Publish(a.events, a.stop)
	go a.Heartbeat(a.ttl, a.stop)

	// Block until we receive a stop signal
	<-a.stop

	a.events.RemoveListener("agent", a.machine)
}

// Stop all async processes the Agent is running
func (a *Agent) Stop() {
	log.V(1).Info("Stopping Agent")
	close(a.stop)
}

// Clear any presence data from the Registry
func (a *Agent) Purge() {
	log.V(1).Info("Removing Agent from Registry")
	bootID := a.machine.State().BootID
	err := a.registry.RemoveMachineState(bootID)
	if err != nil {
		log.Errorf("Failed to remove Machine %s from Registry: %s", bootID, err.Error())
	}

	for _, j := range a.registry.GetAllJobsByMachine(bootID) {
		a.VerifyJob(&j)

		log.V(1).Infof("Clearing JobState(%s) from Registry", j.Name)
		a.registry.RemoveJobState(j.Name)

		// TODO(uwedeportivo): agent placing offer ?
		offer := job.NewOfferFromJob(j, nil)
		log.V(2).Infof("Publishing JobOffer(%s)", offer.Job.Name)
		a.registry.CreateJobOffer(offer)
		log.Infof("Published JobOffer(%s)", offer.Job.Name)
	}
}

// Periodically report to the Registry at an interval equal to
// half of the provided ttl. Stop reporting when the provided
// channel is closed. Failed attempts to report state to the
// Registry are retried twice before moving on to the next
// reporting interval.
func (a *Agent) Heartbeat(ttl time.Duration, stop chan bool) {
	attempt := func(attempts int, f func() error) (err error) {
		if attempts < 1 {
			return fmt.Errorf("attempts argument must be 1 or greater, got %d", attempts)
		}

		// The amount of time the retry mechanism waits after a failed attempt
		// doubles following each failure. This is a simple exponential backoff.
		sleep := time.Second

		for i := 1; i <= attempts; i++ {
			err = f()
			if err == nil || i == attempts {
				break
			}

			sleep = sleep * 2
			log.V(2).Infof("function returned err, retrying in %v: %v", sleep, err)
			time.Sleep(sleep)
		}

		return err
	}

	heartbeat := func() error {
		return a.registry.SetMachineState(a.machine.State(), a.ttl)
	}

	// Explicitly heartbeat immediately to push state to the
	// Registry as quickly as possible
	a.machine.RefreshState()
	if err := attempt(3, heartbeat); err != nil {
		log.Errorf("Failed heartbeat after 3 attempts: %v", err)
	}

	interval := ttl / refreshInterval
	for true {
		select {
		case <-stop:
			log.V(2).Info("MachineHeartbeat exiting due to stop signal")
			return
		case <-time.Tick(interval):
			log.V(2).Info("MachineHeartbeat tick")
			a.machine.RefreshState()
			if err := attempt(3, heartbeat); err != nil {
				log.Errorf("Failed heartbeat after 3 attempts: %v", err)
			}
		}
	}
}

// Instruct the Agent to start the provided Job
func (a *Agent) StartJob(j *job.Job) {
	a.state.TrackJobConflicts(j.Name, j.Payload.Conflicts())

	log.Infof("Starting Job(%s)", j.Name)
	a.systemd.StartJob(j)
}

// Inform the Registry that a Job must be rescheduled
func (a *Agent) RescheduleJob(j *job.Job) {
	log.V(2).Infof("Stopping Job(%s)", j.Name)
	a.registry.UnscheduleJob(j.Name)

	// TODO(uwedeportivo): agent placing offer ?
	offer := job.NewOfferFromJob(*j, nil)
	log.V(2).Infof("Publishing JobOffer(%s)", offer.Job.Name)
	a.registry.CreateJobOffer(offer)
	log.Infof("Published JobOffer(%s)", offer.Job.Name)
}

// Instruct the Agent to stop the provided Job and
// all of its peers
func (a *Agent) StopJob(jobName string) {
	log.Infof("Stopping Job(%s)", jobName)
	a.systemd.StopJob(jobName)
	a.ReportJobState(jobName, nil)

	a.state.Lock()
	reversePeers := a.state.GetJobsByPeer(jobName)
	a.state.DropPeersJob(jobName)
	a.state.DropJobConflicts(jobName)
	a.state.Unlock()

	for _, peer := range reversePeers {
		log.Infof("Stopping Peer(%s) of Job(%s)", peer, jobName)
		a.registry.StopJob(peer)
	}
}

// Persist the state of the given Job into the Registry
func (a *Agent) ReportJobState(jobName string, jobState *job.JobState) {
	if jobState == nil {
		err := a.registry.RemoveJobState(jobName)
		if err != nil {
			log.V(1).Infof("Failed to remove JobState from Registry: %s", jobName, err.Error())
		}
	} else {
		a.registry.SaveJobState(jobName, jobState)
	}
}

// Submit all possible bids for unresolved offers
func (a *Agent) BidForPossibleJobs() {
	a.state.Lock()
	offers := a.state.GetOffersWithoutBids()
	a.state.Unlock()

	log.V(2).Infof("Checking %d unbade offers", len(offers))
	for i, _ := range offers {
		offer := offers[i]
		log.V(2).Infof("Checking ability to run Job(%s)", offer.Job.Name)
		if a.AbleToRun(&offer.Job) {
			log.V(2).Infof("Able to run Job(%s), submitting bid", offer.Job.Name)
			a.Bid(offer.Job.Name)
		} else {
			log.V(2).Infof("Still unable to run Job(%s)", offer.Job.Name)
		}
	}
}

// Submit a bid for the given Job
func (a *Agent) Bid(jobName string) {
	log.Infof("Submitting JobBid for Job(%s)", jobName)

	jb := job.NewBid(jobName, a.machine.State().BootID)
	a.registry.SubmitJobBid(jb)

	a.state.Lock()
	defer a.state.Unlock()

	a.state.TrackBid(jb.JobName)
}

// Instruct the Agent that an offer has been created and must
// be tracked until it is resolved
func (a *Agent) TrackOffer(jo job.JobOffer) {
	a.state.Lock()
	defer a.state.Unlock()

	log.V(2).Infof("Tracking JobOffer(%s)", jo.Job.Name)
	a.state.TrackOffer(jo)

	peers := jo.Job.Payload.Peers()
	log.V(2).Infof("Tracking peers of JobOffer(%s): %v", jo.Job.Name, peers)
	a.state.TrackJobPeers(jo.Job.Name, jo.Job.Payload.Peers())
}

// Instruct the Agent that the given offer has been resolved
// and may be ignored in future conflict calculations
func (a *Agent) OfferResolved(jobName string) {
	a.state.Lock()
	defer a.state.Unlock()

	log.V(2).Infof("Dropping JobOffer(%s)", jobName)
	a.state.DropOffer(jobName)

	a.state.DropBid(jobName)
}

// Pull a Job and its payload from the Registry
func (a *Agent) FetchJob(jobName string) *job.Job {
	log.V(1).Infof("Fetching Job(%s) from Registry", jobName)
	j := a.registry.GetJob(jobName)
	if j == nil {
		log.V(1).Infof("Job not found in Registry")
		return nil
	}
	return j
}

// Verify a Job through SignatureSet
func (a *Agent) VerifyJob(j *job.Job) bool {
	if a.verifier == nil {
		return true
	}

	payload := j.Payload
	s := a.registry.GetSignatureSetOfPayload(payload.Name)
	ok, err := a.verifier.VerifyPayload(payload, s)
	if !ok || err != nil {
		log.V(1).Infof("Payload(%s) doesn't fit its signature: %v", payload.Name, err)
		return false
	}
	return true
}

// Submit all possible bids for known peers of the provided job
func (a *Agent) BidForPossiblePeers(jobName string) {
	a.state.Lock()
	peers := a.state.GetJobsByPeer(jobName)
	a.state.Unlock()

	for _, peer := range peers {
		log.V(1).Infof("Found unresolved offer for Peer(%s) of Job(%s)", peer, jobName)

		peerJob := a.FetchJob(peer)
		if peerJob != nil && a.AbleToRun(peerJob) {
			log.Infof("Submitting bid for Peer(%s) of Job(%s)", peer, jobName)
			a.Bid(peer)
		} else {
			log.V(1).Infof("Unable to bid for Peer(%s) of Job(%s)", peer, jobName)
		}
	}
}

// Determine if the Agent can run the provided Job
func (a *Agent) AbleToRun(j *job.Job) bool {
	if !a.VerifyJob(j) {
		log.V(1).Infof("Failed to verify Job(%s)", j.Name)
		return false
	}

	requirements := j.Requirements()
	if len(requirements) == 0 {
		log.V(1).Infof("Job(%s) has no requirements", j.Name)
		return true
	}

	if log.V(1) {
		var reqString string
		for key, slice := range requirements {
			reqString += fmt.Sprintf("%s = [", key)
			for _, val := range slice {
				reqString += fmt.Sprintf("%s, ", val)
			}
			reqString += fmt.Sprint("] ")
		}

		log.Infof("Job(%s) has requirements: %s", j.Name, reqString)
	}

	metadata := extractMachineMetadata(requirements)
	log.V(1).Infof("Job(%s) requires machine metadata: %v", j.Name, metadata)
	if !a.machine.HasMetadata(metadata) {
		log.V(1).Infof("Unable to run Job(%s), local Machine metadata insufficient", j.Name)
		return false
	}

	bootID, ok := requirements[unit.FleetXConditionMachineBootID]
	if ok && len(bootID) > 0 && !a.machine.State().MatchBootID(bootID[0]) {
		log.V(1).Infof("Agent does not pass MachineBootID condition for Job(%s)", j.Name)
		return false
	}

	peers := j.Payload.Peers()
	if len(peers) > 0 {
		log.V(1).Infof("Asserting required Peers %v of Job(%s) are scheduled locally", peers, j.Name)
		for _, peer := range peers {
			if !a.peerScheduledHere(j.Name, peer) {
				log.V(1).Infof("Required Peer(%s) of Job(%s) is not scheduled locally", peer, j.Name)
				return false
			}
		}
	} else {
		log.V(2).Infof("Job(%s) has no peers to worry about", j.Name)
	}

	if conflicted, conflictedJobName := a.state.HasConflict(j.Name, j.Payload.Conflicts()); conflicted {
		log.V(1).Infof("Job(%s) has conflict with Job(%s)", j.Name, conflictedJobName)
		return false
	}

	return true
}

// Return all machine-related metadata from a job requirements map
func extractMachineMetadata(requirements map[string][]string) map[string][]string {
	metadata := make(map[string][]string)

	for key, values := range requirements {
		if !strings.HasPrefix(key, "MachineMetadata") {
			continue
		}

		// Strip off leading 'MachineMetadata'
		key = key[15:]

		if len(values) == 0 {
			log.V(2).Infof("Machine metadata requirement %s provided no values, ignoring.", key)
			continue
		}

		metadata[key] = values
	}

	return metadata
}

// Determine if all necessary peers of a Job are scheduled to this Agent
func (a *Agent) peerScheduledHere(jobName, peerName string) bool {
	log.V(1).Infof("Looking for target of Peer(%s)", peerName)

	//FIXME: ideally the machine would use its own knowledge rather than calling GetJobTarget
	if tgt := a.registry.GetJobTarget(peerName); tgt == nil || tgt.BootID != a.machine.State().BootID {
		log.V(1).Infof("Peer(%s) of Job(%s) not scheduled here", peerName, jobName)
		return false
	}

	log.V(1).Infof("Peer(%s) of Job(%s) scheduled here", peerName, jobName)
	return true
}

func (a *Agent) UnresolvedJobOffers() []job.JobOffer {
	return a.registry.UnresolvedJobOffers()
}
<commit_msg>fix(agent): Correct reference to agent ttl<commit_after>package agent

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

	log "github.com/coreos/fleet/third_party/github.com/golang/glog"

	"github.com/coreos/fleet/event"
	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/registry"
	"github.com/coreos/fleet/sign"
	"github.com/coreos/fleet/systemd"
	"github.com/coreos/fleet/unit"
)

const (
	// TTL to use with all state pushed to Registry
	DefaultTTL = "30s"

	// Refresh TTLs at 1/2 the TTL length
	refreshInterval = 2
)

// The Agent owns all of the coordination between the Registry, the local
// Machine, and the local SystemdManager.
type Agent struct {
	registry      *registry.Registry
	events        *event.EventBus
	machine       *machine.Machine
	ttl           time.Duration
	systemdPrefix string
	// verifier is used to verify job payload. A nil one implies that
	// all payloads are accepted.
	verifier *sign.SignatureVerifier

	state   *AgentState
	systemd *systemd.SystemdManager

	// channel used to shutdown any open connections/channels the Agent holds
	stop chan bool
}

func New(registry *registry.Registry, events *event.EventBus, machine *machine.Machine, ttl, unitPrefix string, verifier *sign.SignatureVerifier) (*Agent, error) {
	ttldur, err := time.ParseDuration(ttl)
	if err != nil {
		return nil, err
	}

	state := NewState()
	mgr := systemd.NewSystemdManager(machine, unitPrefix)

	return &Agent{registry, events, machine, ttldur, unitPrefix, verifier, state, mgr, nil}, nil
}

// Access Agent's machine field
func (a *Agent) Machine() *machine.Machine {
	return a.machine
}

func (a *Agent) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct{ State *AgentState }{State: a.state})
}

// Trigger all async processes the Agent intends to run
func (a *Agent) Run() {
	a.stop = make(chan bool)

	handler := NewEventHandler(a)
	a.events.AddListener("agent", a.machine, handler)

	go a.systemd.Publish(a.events, a.stop)
	go a.Heartbeat(a.ttl, a.stop)

	// Block until we receive a stop signal
	<-a.stop

	a.events.RemoveListener("agent", a.machine)
}

// Stop all async processes the Agent is running
func (a *Agent) Stop() {
	log.V(1).Info("Stopping Agent")
	close(a.stop)
}

// Clear any presence data from the Registry
func (a *Agent) Purge() {
	log.V(1).Info("Removing Agent from Registry")
	bootID := a.machine.State().BootID
	err := a.registry.RemoveMachineState(bootID)
	if err != nil {
		log.Errorf("Failed to remove Machine %s from Registry: %s", bootID, err.Error())
	}

	for _, j := range a.registry.GetAllJobsByMachine(bootID) {
		a.VerifyJob(&j)

		log.V(1).Infof("Clearing JobState(%s) from Registry", j.Name)
		a.registry.RemoveJobState(j.Name)

		// TODO(uwedeportivo): agent placing offer ?
		offer := job.NewOfferFromJob(j, nil)
		log.V(2).Infof("Publishing JobOffer(%s)", offer.Job.Name)
		a.registry.CreateJobOffer(offer)
		log.Infof("Published JobOffer(%s)", offer.Job.Name)
	}
}

// Periodically report to the Registry at an interval equal to
// half of the provided ttl. Stop reporting when the provided
// channel is closed. Failed attempts to report state to the
// Registry are retried twice before moving on to the next
// reporting interval.
func (a *Agent) Heartbeat(ttl time.Duration, stop chan bool) {
	attempt := func(attempts int, f func() error) (err error) {
		if attempts < 1 {
			return fmt.Errorf("attempts argument must be 1 or greater, got %d", attempts)
		}

		// The amount of time the retry mechanism waits after a failed attempt
		// doubles following each failure. This is a simple exponential backoff.
		sleep := time.Second

		for i := 1; i <= attempts; i++ {
			err = f()
			if err == nil || i == attempts {
				break
			}

			sleep = sleep * 2
			log.V(2).Infof("function returned err, retrying in %v: %v", sleep, err)
			time.Sleep(sleep)
		}

		return err
	}

	heartbeat := func() error {
		return a.registry.SetMachineState(a.machine.State(), ttl)
	}

	// Explicitly heartbeat immediately to push state to the
	// Registry as quickly as possible
	a.machine.RefreshState()
	if err := attempt(3, heartbeat); err != nil {
		log.Errorf("Failed heartbeat after 3 attempts: %v", err)
	}

	interval := ttl / refreshInterval
	for true {
		select {
		case <-stop:
			log.V(2).Info("MachineHeartbeat exiting due to stop signal")
			return
		case <-time.Tick(interval):
			log.V(2).Info("MachineHeartbeat tick")
			a.machine.RefreshState()
			if err := attempt(3, heartbeat); err != nil {
				log.Errorf("Failed heartbeat after 3 attempts: %v", err)
			}
		}
	}
}

// Instruct the Agent to start the provided Job
func (a *Agent) StartJob(j *job.Job) {
	a.state.TrackJobConflicts(j.Name, j.Payload.Conflicts())

	log.Infof("Starting Job(%s)", j.Name)
	a.systemd.StartJob(j)
}

// Inform the Registry that a Job must be rescheduled
func (a *Agent) RescheduleJob(j *job.Job) {
	log.V(2).Infof("Stopping Job(%s)", j.Name)
	a.registry.UnscheduleJob(j.Name)

	// TODO(uwedeportivo): agent placing offer ?
	offer := job.NewOfferFromJob(*j, nil)
	log.V(2).Infof("Publishing JobOffer(%s)", offer.Job.Name)
	a.registry.CreateJobOffer(offer)
	log.Infof("Published JobOffer(%s)", offer.Job.Name)
}

// Instruct the Agent to stop the provided Job and
// all of its peers
func (a *Agent) StopJob(jobName string) {
	log.Infof("Stopping Job(%s)", jobName)
	a.systemd.StopJob(jobName)
	a.ReportJobState(jobName, nil)

	a.state.Lock()
	reversePeers := a.state.GetJobsByPeer(jobName)
	a.state.DropPeersJob(jobName)
	a.state.DropJobConflicts(jobName)
	a.state.Unlock()

	for _, peer := range reversePeers {
		log.Infof("Stopping Peer(%s) of Job(%s)", peer, jobName)
		a.registry.StopJob(peer)
	}
}

// Persist the state of the given Job into the Registry
func (a *Agent) ReportJobState(jobName string, jobState *job.JobState) {
	if jobState == nil {
		err := a.registry.RemoveJobState(jobName)
		if err != nil {
			log.V(1).Infof("Failed to remove JobState from Registry: %s", jobName, err.Error())
		}
	} else {
		a.registry.SaveJobState(jobName, jobState)
	}
}

// Submit all possible bids for unresolved offers
func (a *Agent) BidForPossibleJobs() {
	a.state.Lock()
	offers := a.state.GetOffersWithoutBids()
	a.state.Unlock()

	log.V(2).Infof("Checking %d unbade offers", len(offers))
	for i, _ := range offers {
		offer := offers[i]
		log.V(2).Infof("Checking ability to run Job(%s)", offer.Job.Name)
		if a.AbleToRun(&offer.Job) {
			log.V(2).Infof("Able to run Job(%s), submitting bid", offer.Job.Name)
			a.Bid(offer.Job.Name)
		} else {
			log.V(2).Infof("Still unable to run Job(%s)", offer.Job.Name)
		}
	}
}

// Submit a bid for the given Job
func (a *Agent) Bid(jobName string) {
	log.Infof("Submitting JobBid for Job(%s)", jobName)

	jb := job.NewBid(jobName, a.machine.State().BootID)
	a.registry.SubmitJobBid(jb)

	a.state.Lock()
	defer a.state.Unlock()

	a.state.TrackBid(jb.JobName)
}

// Instruct the Agent that an offer has been created and must
// be tracked until it is resolved
func (a *Agent) TrackOffer(jo job.JobOffer) {
	a.state.Lock()
	defer a.state.Unlock()

	log.V(2).Infof("Tracking JobOffer(%s)", jo.Job.Name)
	a.state.TrackOffer(jo)

	peers := jo.Job.Payload.Peers()
	log.V(2).Infof("Tracking peers of JobOffer(%s): %v", jo.Job.Name, peers)
	a.state.TrackJobPeers(jo.Job.Name, jo.Job.Payload.Peers())
}

// Instruct the Agent that the given offer has been resolved
// and may be ignored in future conflict calculations
func (a *Agent) OfferResolved(jobName string) {
	a.state.Lock()
	defer a.state.Unlock()

	log.V(2).Infof("Dropping JobOffer(%s)", jobName)
	a.state.DropOffer(jobName)

	a.state.DropBid(jobName)
}

// Pull a Job and its payload from the Registry
func (a *Agent) FetchJob(jobName string) *job.Job {
	log.V(1).Infof("Fetching Job(%s) from Registry", jobName)
	j := a.registry.GetJob(jobName)
	if j == nil {
		log.V(1).Infof("Job not found in Registry")
		return nil
	}
	return j
}

// Verify a Job through SignatureSet
func (a *Agent) VerifyJob(j *job.Job) bool {
	if a.verifier == nil {
		return true
	}

	payload := j.Payload
	s := a.registry.GetSignatureSetOfPayload(payload.Name)
	ok, err := a.verifier.VerifyPayload(payload, s)
	if !ok || err != nil {
		log.V(1).Infof("Payload(%s) doesn't fit its signature: %v", payload.Name, err)
		return false
	}
	return true
}

// Submit all possible bids for known peers of the provided job
func (a *Agent) BidForPossiblePeers(jobName string) {
	a.state.Lock()
	peers := a.state.GetJobsByPeer(jobName)
	a.state.Unlock()

	for _, peer := range peers {
		log.V(1).Infof("Found unresolved offer for Peer(%s) of Job(%s)", peer, jobName)

		peerJob := a.FetchJob(peer)
		if peerJob != nil && a.AbleToRun(peerJob) {
			log.Infof("Submitting bid for Peer(%s) of Job(%s)", peer, jobName)
			a.Bid(peer)
		} else {
			log.V(1).Infof("Unable to bid for Peer(%s) of Job(%s)", peer, jobName)
		}
	}
}

// Determine if the Agent can run the provided Job
func (a *Agent) AbleToRun(j *job.Job) bool {
	if !a.VerifyJob(j) {
		log.V(1).Infof("Failed to verify Job(%s)", j.Name)
		return false
	}

	requirements := j.Requirements()
	if len(requirements) == 0 {
		log.V(1).Infof("Job(%s) has no requirements", j.Name)
		return true
	}

	if log.V(1) {
		var reqString string
		for key, slice := range requirements {
			reqString += fmt.Sprintf("%s = [", key)
			for _, val := range slice {
				reqString += fmt.Sprintf("%s, ", val)
			}
			reqString += fmt.Sprint("] ")
		}

		log.Infof("Job(%s) has requirements: %s", j.Name, reqString)
	}

	metadata := extractMachineMetadata(requirements)
	log.V(1).Infof("Job(%s) requires machine metadata: %v", j.Name, metadata)
	if !a.machine.HasMetadata(metadata) {
		log.V(1).Infof("Unable to run Job(%s), local Machine metadata insufficient", j.Name)
		return false
	}

	bootID, ok := requirements[unit.FleetXConditionMachineBootID]
	if ok && len(bootID) > 0 && !a.machine.State().MatchBootID(bootID[0]) {
		log.V(1).Infof("Agent does not pass MachineBootID condition for Job(%s)", j.Name)
		return false
	}

	peers := j.Payload.Peers()
	if len(peers) > 0 {
		log.V(1).Infof("Asserting required Peers %v of Job(%s) are scheduled locally", peers, j.Name)
		for _, peer := range peers {
			if !a.peerScheduledHere(j.Name, peer) {
				log.V(1).Infof("Required Peer(%s) of Job(%s) is not scheduled locally", peer, j.Name)
				return false
			}
		}
	} else {
		log.V(2).Infof("Job(%s) has no peers to worry about", j.Name)
	}

	if conflicted, conflictedJobName := a.state.HasConflict(j.Name, j.Payload.Conflicts()); conflicted {
		log.V(1).Infof("Job(%s) has conflict with Job(%s)", j.Name, conflictedJobName)
		return false
	}

	return true
}

// Return all machine-related metadata from a job requirements map
func extractMachineMetadata(requirements map[string][]string) map[string][]string {
	metadata := make(map[string][]string)

	for key, values := range requirements {
		if !strings.HasPrefix(key, "MachineMetadata") {
			continue
		}

		// Strip off leading 'MachineMetadata'
		key = key[15:]

		if len(values) == 0 {
			log.V(2).Infof("Machine metadata requirement %s provided no values, ignoring.", key)
			continue
		}

		metadata[key] = values
	}

	return metadata
}

// Determine if all necessary peers of a Job are scheduled to this Agent
func (a *Agent) peerScheduledHere(jobName, peerName string) bool {
	log.V(1).Infof("Looking for target of Peer(%s)", peerName)

	//FIXME: ideally the machine would use its own knowledge rather than calling GetJobTarget
	if tgt := a.registry.GetJobTarget(peerName); tgt == nil || tgt.BootID != a.machine.State().BootID {
		log.V(1).Infof("Peer(%s) of Job(%s) not scheduled here", peerName, jobName)
		return false
	}

	log.V(1).Infof("Peer(%s) of Job(%s) scheduled here", peerName, jobName)
	return true
}

func (a *Agent) UnresolvedJobOffers() []job.JobOffer {
	return a.registry.UnresolvedJobOffers()
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Pani Networks
// 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.

// mocks.go contins Helper structure with interfaces which which is used
// to interact with operation system.
//
// - OS interface used to access filesystem, write and read files.
// - Executable interface is used to execute commands in operation system.
//
// Both interfaces has default and fake implementations, default implementation
// will usually just proxy calls to standard library while test implemetation
// will allow mocking all interactions.

package agent

import (
	"fmt"
	"github.com/romana/core/common"

	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"strings"
	"sync"
)

// TODO There is a tradeoff, either use global variable for provider
// or pass provider down to each method.
// Passing down to each method is more explicit which is good,
// but pollutes methods' signatures too much. Need to have a discussion.

// Helper groups testable implementations
// of standard library functions.
type Helper struct {
	Executor                   Executable
	OS                         OS
	Agent                      *Agent //access field for Agent
	ensureRouteToEndpointMutex *sync.Mutex
	ensureLineMutex            *sync.Mutex
	ensureInterHostRoutesMutex *sync.Mutex
}

// mockAgent creates the agent with the configuration
// needed for tests without the need to go through
// configuration files.
func mockAgent() Agent {

	host0 := common.HostMessage{Ip: "172.17.0.1", RomanaIp: "127.0.0.1/8"}

	romanaIP, romanaNet, _ := net.ParseCIDR(host0.RomanaIp)
	networkConfig := &NetworkConfig{}
	networkConfig.currentHostIP = net.ParseIP(host0.Ip)
	networkConfig.currentHostGW = romanaIP
	networkConfig.currentHostGWNet = *romanaNet
	networkConfig.currentHostGWNetSize, _ = romanaNet.Mask.Size()
	networkConfig.currentHostIndex = 0

	host1 := common.HostMessage{Ip: "192.168.0.12", RomanaIp: "10.65.0.0/16"}
	networkConfig.hosts = []common.HostMessage{host0, host1}

	dc := common.Datacenter{}
	dc.Cidr = "10.0.0.0/8"
	dc.PortBits = 8
	dc.TenantBits = 4
	dc.SegmentBits = 4
	dc.EndpointSpaceBits = 0
	dc.EndpointBits = 8

	networkConfig.dc = dc

	agent := &Agent{networkConfig: networkConfig}
	helper := NewAgentHelper(agent)
	agent.Helper = &helper

	return *agent
}

// Executable is an interface that mocks exec.Command().Output()
type Executable interface {
	Exec(cmd string, args []string) ([]byte, error)
}

// DefaultExecutor is a default implementation of Executable that passes
// back to standard library.
type DefaultExecutor struct{}

// Exec proxies all requests to exec.Command()
// Used to support unit testing.
func (DefaultExecutor) Exec(cmd string, args []string) ([]byte, error) {
	log.Printf("Helper.Executor: executing command: %s %s", cmd, args)
	out, err := exec.Command(cmd, args...).Output()
	return out, err
}

// OS interface mocks standard lib os.
type OS interface {
	open(name string) (OSFile, error)
	appendFile(name string) (OSFile, error)
	createIfMissing(name string) error
}

// OSFile interface mocks os.File
type OSFile interface {
	io.Reader
	io.Writer
	io.Closer
}

// DefaultOS is a default implementation of OS interface
// which proxies everything to standard lib.
type DefaultOS struct {
}

// open is a direct proxy to os.Open
func (DefaultOS) open(name string) (OSFile, error) {
	f, err := os.Open(name)
	return f, err
}

// appendFile returns a file opened for write
// with cursor positioned at the end of file.
func (DefaultOS) appendFile(name string) (OSFile, error) {
	file, err := os.OpenFile(name, os.O_APPEND|os.O_WRONLY, 0600)
	return file, err
}

// createIfMissing tries create file if it's not there yet,
// otherwise no op.
func (DefaultOS) createIfMissing(name string) error {
	file, err := os.OpenFile(name, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	file.Close()
	return nil
}

// FakeFile implements OSFile.
type FakeFile struct {
	io.Reader
	content string
}

// Close is a no op to satisfy OSFile.
func (f FakeFile) Close() error {
	return nil
}

// Write is a method of FakeFile that records all data it receives.
func (f *FakeFile) Write(p []byte) (ret int, err error) {
	f.content = fmt.Sprintf("%s%s", f.content, string(p))
	ret = len(p)
	return ret, nil
}

// FakeOS implements OS.
type FakeOS struct {
	fakeData string
	fakeFile *FakeFile
}

// open returns a FakeFile stuffed with fake data
func (o FakeOS) open(name string) (OSFile, error) {
	fake := FakeFile{strings.NewReader(o.fakeData), ""}
	return &fake, nil
}

// appendFile returns a FakeFile implementation
// that will record any data it receives for later analisis.
func (o *FakeOS) appendFile(name string) (OSFile, error) {
	fake := FakeFile{strings.NewReader(o.fakeData), ""}
	o.fakeFile = &fake
	return &fake, nil
}

// createIfMissing No op in tests.
func (o *FakeOS) createIfMissing(name string) error {
	return nil
}

// FakeExecutor implements Executable
// stores faked Output, Error and commands recorded by Exec.
type FakeExecutor struct {
	Output   []byte
	Error    error
	Commands *string
}

// Exec is a method of fake executor that will record all incoming commands
// and use faked Output and Error.
func (x *FakeExecutor) Exec(cmd string, args []string) ([]byte, error) {
	var c string
	if x.Commands == nil {
		c = fmt.Sprintf("%s %s", cmd, strings.Join(args, " "))
	} else {
		c = fmt.Sprintf("%s\n%s %s", *x.Commands, cmd, strings.Join(args, " "))
	}
	x.Commands = &c
	return x.Output, x.Error
}
<commit_msg>Update mock. More updates required to pass tests.<commit_after>// Copyright (c) 2016 Pani Networks
// 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.

// mocks.go contins Helper structure with interfaces which which is used
// to interact with operation system.
//
// - OS interface used to access filesystem, write and read files.
// - Executable interface is used to execute commands in operation system.
//
// Both interfaces has default and fake implementations, default implementation
// will usually just proxy calls to standard library while test implemetation
// will allow mocking all interactions.

package agent

import (
	"fmt"
	"github.com/romana/core/common"

	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"strings"
	"sync"
)

// TODO There is a tradeoff, either use global variable for provider
// or pass provider down to each method.
// Passing down to each method is more explicit which is good,
// but pollutes methods' signatures too much. Need to have a discussion.

// Helper groups testable implementations
// of standard library functions.
type Helper struct {
	Executor                   Executable
	OS                         OS
	Agent                      *Agent //access field for Agent
	ensureRouteToEndpointMutex *sync.Mutex
	ensureLineMutex            *sync.Mutex
	ensureInterHostRoutesMutex *sync.Mutex
}

// mockAgent creates the agent with the configuration
// needed for tests without the need to go through
// configuration files.
func mockAgent() Agent {

	host0 := common.HostMessage{Ip: "172.17.0.1", RomanaIp: "127.0.0.1/8"}

	// romanaIP, romanaNet, _ := net.ParseCIDR(host0.RomanaIp)

	networkConfig := &NetworkConfig{}
	networkConfig.romanaIP = net.ParseIP(host0.Ip)

	host1 := common.HostMessage{Ip: "192.168.0.12", RomanaIp: "10.65.0.0/16"}
	networkConfig.otherHosts = []common.HostMessage{host1}

	dc := common.Datacenter{}
	dc.Cidr = "10.0.0.0/8"
	dc.PortBits = 8
	dc.TenantBits = 4
	dc.SegmentBits = 4
	dc.EndpointSpaceBits = 0
	dc.EndpointBits = 8

	networkConfig.dc = dc
	agent := &Agent{networkConfig: networkConfig}
	helper := NewAgentHelper(agent)
	agent.Helper = &helper

	return *agent
}

// Executable is an interface that mocks exec.Command().Output()
type Executable interface {
	Exec(cmd string, args []string) ([]byte, error)
}

// DefaultExecutor is a default implementation of Executable that passes
// back to standard library.
type DefaultExecutor struct{}

// Exec proxies all requests to exec.Command()
// Used to support unit testing.
func (DefaultExecutor) Exec(cmd string, args []string) ([]byte, error) {
	log.Printf("Helper.Executor: executing command: %s %s", cmd, args)
	out, err := exec.Command(cmd, args...).Output()
	return out, err
}

// OS interface mocks standard lib os.
type OS interface {
	open(name string) (OSFile, error)
	appendFile(name string) (OSFile, error)
	createIfMissing(name string) error
}

// OSFile interface mocks os.File
type OSFile interface {
	io.Reader
	io.Writer
	io.Closer
}

// DefaultOS is a default implementation of OS interface
// which proxies everything to standard lib.
type DefaultOS struct {
}

// open is a direct proxy to os.Open
func (DefaultOS) open(name string) (OSFile, error) {
	f, err := os.Open(name)
	return f, err
}

// appendFile returns a file opened for write
// with cursor positioned at the end of file.
func (DefaultOS) appendFile(name string) (OSFile, error) {
	file, err := os.OpenFile(name, os.O_APPEND|os.O_WRONLY, 0600)
	return file, err
}

// createIfMissing tries create file if it's not there yet,
// otherwise no op.
func (DefaultOS) createIfMissing(name string) error {
	file, err := os.OpenFile(name, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	file.Close()
	return nil
}

// FakeFile implements OSFile.
type FakeFile struct {
	io.Reader
	content string
}

// Close is a no op to satisfy OSFile.
func (f FakeFile) Close() error {
	return nil
}

// Write is a method of FakeFile that records all data it receives.
func (f *FakeFile) Write(p []byte) (ret int, err error) {
	f.content = fmt.Sprintf("%s%s", f.content, string(p))
	ret = len(p)
	return ret, nil
}

// FakeOS implements OS.
type FakeOS struct {
	fakeData string
	fakeFile *FakeFile
}

// open returns a FakeFile stuffed with fake data
func (o FakeOS) open(name string) (OSFile, error) {
	fake := FakeFile{strings.NewReader(o.fakeData), ""}
	return &fake, nil
}

// appendFile returns a FakeFile implementation
// that will record any data it receives for later analisis.
func (o *FakeOS) appendFile(name string) (OSFile, error) {
	fake := FakeFile{strings.NewReader(o.fakeData), ""}
	o.fakeFile = &fake
	return &fake, nil
}

// createIfMissing No op in tests.
func (o *FakeOS) createIfMissing(name string) error {
	return nil
}

// FakeExecutor implements Executable
// stores faked Output, Error and commands recorded by Exec.
type FakeExecutor struct {
	Output   []byte
	Error    error
	Commands *string
}

// Exec is a method of fake executor that will record all incoming commands
// and use faked Output and Error.
func (x *FakeExecutor) Exec(cmd string, args []string) ([]byte, error) {
	var c string
	if x.Commands == nil {
		c = fmt.Sprintf("%s %s", cmd, strings.Join(args, " "))
	} else {
		c = fmt.Sprintf("%s\n%s %s", *x.Commands, cmd, strings.Join(args, " "))
	}
	x.Commands = &c
	return x.Output, x.Error
}
<|endoftext|>


<commit_before>package godingtalk

import (
    "fmt"
    "net/url"
)

type User struct {
    OAPIResponse
    Userid string
    Name string
    Mobile string
    Tel string
    Remark string
    Order int
    IsAdmin bool
    IsBoss bool
    IsLeader bool
	IsSys bool `json:"is_sys"`
	SysLevel int `json:"sys_level"`
    Active bool
    Department []int
    Position string
    Email string
    Avatar string
    Extattr interface{}
}

type UserList struct {
    OAPIResponse
    HasMore bool
    Userlist []User
}

type Department struct {
    OAPIResponse
    Id int
    Name string
    ParentId int
    Order int
    DeptPerimits string
    UserPerimits string
    OuterDept bool
    OuterPermitDepts string
    OuterPermitUsers string
    OrgDeptOwner string
    DeptManagerUseridList string
}

type DepartmentList struct {
    OAPIResponse
    Departments []Department `json:"department"`
}

// DepartmentList is 获取部门列表
func (c *DingTalkClient) DepartmentList() (DepartmentList, error) {
    var data DepartmentList
    err := c.httpRPC("department/list", nil, nil, &data)
    return data, err
}

//DepartmentDetail is 获取部门详情
func (c *DingTalkClient) DepartmentDetail(id int) (Department, error) {
    var data Department
    params := url.Values{}
    params.Add("id", fmt.Sprintf("%d", id))
    err :=c.httpRPC("department/get", params, nil, &data)
    return data, err
}

//UserList is 获取部门成员
func (c *DingTalkClient) UserList(departmentID int) (UserList, error) {
    var data UserList
    params := url.Values{}
    params.Add("department_id", fmt.Sprintf("%d", departmentID))
    err :=c.httpRPC("user/simplelist", params, nil, &data)
    return data, err
}

//CreateChat is
func (c *DingTalkClient) CreateChat(name string, owner string, useridlist []string) (string, error) {
    var data struct {
        OAPIResponse
        Chatid string
    }
    request := map[string]interface{} {
        "name":name,
        "owner":owner,
        "useridlist":useridlist,
    }
    err :=c.httpRPC("chat/create", nil, request, &data)
    return data.Chatid, err
}

//UserInfoByCode 校验免登录码并换取用户身份
func (c *DingTalkClient) UserInfoByCode(code string) (User, error) {
    var data User
    params := url.Values{}
    params.Add("code", code)
    err :=c.httpRPC("user/getuserinfo", params, nil, &data)
    return data, err
}<commit_msg>Replace user/simplelist with user/list in UserList<commit_after>package godingtalk

import (
    "fmt"
    "net/url"
)

type User struct {
    OAPIResponse
    Userid string
    Name string
    Mobile string
    Tel string
    Remark string
    Order int
    IsAdmin bool
    IsBoss bool
    IsLeader bool
	IsSys bool `json:"is_sys"`
	SysLevel int `json:"sys_level"`
    Active bool
    Department []int
    Position string
    Email string
    Avatar string
    Extattr interface{}
}

type UserList struct {
    OAPIResponse
    HasMore bool
    Userlist []User
}

type Department struct {
    OAPIResponse
    Id int
    Name string
    ParentId int
    Order int
    DeptPerimits string
    UserPerimits string
    OuterDept bool
    OuterPermitDepts string
    OuterPermitUsers string
    OrgDeptOwner string
    DeptManagerUseridList string
}

type DepartmentList struct {
    OAPIResponse
    Departments []Department `json:"department"`
}

// DepartmentList is 获取部门列表
func (c *DingTalkClient) DepartmentList() (DepartmentList, error) {
    var data DepartmentList
    err := c.httpRPC("department/list", nil, nil, &data)
    return data, err
}

//DepartmentDetail is 获取部门详情
func (c *DingTalkClient) DepartmentDetail(id int) (Department, error) {
    var data Department
    params := url.Values{}
    params.Add("id", fmt.Sprintf("%d", id))
    err :=c.httpRPC("department/get", params, nil, &data)
    return data, err
}

//UserList is 获取部门成员
func (c *DingTalkClient) UserList(departmentID int) (UserList, error) {
    var data UserList
    params := url.Values{}
    params.Add("department_id", fmt.Sprintf("%d", departmentID))
    err :=c.httpRPC("user/list", params, nil, &data)
    return data, err
}

//CreateChat is
func (c *DingTalkClient) CreateChat(name string, owner string, useridlist []string) (string, error) {
    var data struct {
        OAPIResponse
        Chatid string
    }
    request := map[string]interface{} {
        "name":name,
        "owner":owner,
        "useridlist":useridlist,
    }
    err :=c.httpRPC("chat/create", nil, request, &data)
    return data.Chatid, err
}

//UserInfoByCode 校验免登录码并换取用户身份
func (c *DingTalkClient) UserInfoByCode(code string) (User, error) {
    var data User
    params := url.Values{}
    params.Add("code", code)
    err :=c.httpRPC("user/getuserinfo", params, nil, &data)
    return data, err
}<|endoftext|>


<commit_before>package config

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/omakoto/go-common/src/fileutils"
	"github.com/omakoto/go-common/src/utils"
	"github.com/omakoto/zenlog/zenlog/envs"
	"github.com/omakoto/zenlog/zenlog/util"
	"runtime"
)

var (
	isLogger    = false
	isLoggerSet = false
	config      *Config
)

func SetIsLogger(logger bool) {
	if isLoggerSet {
		util.Fatalf("SetIsLogger already set")
	}
	isLogger = logger
	isLoggerSet = true
}

// Config represents configuration parameters, read from ~/.zenlog.toml and overridden with the environmental variables.
type Config struct {
	LogDir              string `toml:"ZENLOG_DIR"`
	TempDir             string `toml:"ZENLOG_TEMP"`
	SourceDir           string `toml:"ZENLOG_SOURCE_DIR"`
	StartCommand        string `toml:"ZENLOG_START_COMMAND"`
	PrefixCommands      string `toml:"ZENLOG_PREFIX_COMMANDS"`
	AlwaysNoLogCommands string `toml:"ZENLOG_ALWAYS_NO_LOG_COMMANDS"`
	AutoFlush           bool   `toml:"ZENLOG_AUTO_FLUSH"`

	UseExperimentalCommandParser bool   `toml:"ZENLOG_USE_EXPERIMENTAL_COMMAND_PARSER"`
	CommandSplitter              string `toml:"ZENLOG_COMMAND_SPLITTER"`
	CommentSplitter              string `toml:"ZENLOG_COMMENT_SPLITTER"`

	Maxproc                 int `toml:"ZENLOG_MAXPROC"`
	CriticalCrashMaxSeconds int `toml:"ZENLOG_CRITICAL_CRASH_MAX_SECONDS"`

	UseSplice bool `toml:"ZENLOG_USE_SPLICE"`

	ZenlogPid int

	OuterTty  string
	LoggerIn  string
	LoggerOut string
}

// If a string is empty, get from a given environmental variabne.
func overwriteWithEnviron(to *string, envKey string, def string) {
	if val := os.Getenv(envKey); val != "" {
		*to = val
	} else {
		*to = os.ExpandEnv(*to)
	}
	if *to == "" {
		*to = def
	}
}

func overwriteBoolWithEnviron(to *bool, envKey string) {
	if val := os.Getenv(envKey); val != "" {
		*to = (val == "1")
	}
}

// Make sure a string ends with a slash.
func ensureSlash(v *string) {
	if strings.HasSuffix(*v, "/") {
		return
	}
	*v = *v + "/"
}

// InitConfigForLogger returns a Config for a new session, loading from ~/.zenlog.toml and the environmental variables.
func InitConfigForLogger() *Config {
	file := os.Getenv(envs.ZenlogConf)
	if file == "" {
		file = os.ExpandEnv("$HOME/.zenlog.toml")
	}
	var c Config

	c.UseExperimentalCommandParser = true // Default to true.
	c.UseSplice = true

	util.Debugf("config=%s", file)

	data, err := ioutil.ReadFile(file)
	if err == nil {
		if _, err := toml.Decode(string(data), &c); err != nil {
			util.Fatalf("Unable to load %s: %s", file, err)
		}
	} else if os.IsNotExist(err) {
		util.Warn(err, "%s doesn't exist; using the default instead", file)
	}

	overwriteWithEnviron(&c.StartCommand, "ZENLOG_START_COMMAND", "")
	overwriteWithEnviron(&c.LogDir, envs.ZenlogDir, os.ExpandEnv("$HOME/zenlog/"))
	overwriteWithEnviron(&c.SourceDir, envs.ZenlogSourceDir, "")
	overwriteWithEnviron(&c.PrefixCommands, "ZENLOG_PREFIX_COMMANDS", `(?:command|builtin|time|sudo|[a-zA-Z0-9_]+\=.*)`)
	overwriteWithEnviron(&c.AlwaysNoLogCommands, "ZENLOG_ALWAYS_NO_LOG_COMMANDS", `(?:vi|vim|man|nano|pico|emacs|zenlog.*)`)

	overwriteWithEnviron(&c.CommandSplitter, "ZENLOG_COMMAND_SPLITTER", "")
	overwriteWithEnviron(&c.CommentSplitter, "ZENLOG_COMMENT_SPLITTER", "")

	overwriteWithEnviron(&c.TempDir, envs.ZenlogTemp, "")

	overwriteBoolWithEnviron(&c.AutoFlush, envs.ZenlogAutoFlush)
	overwriteBoolWithEnviron(&c.UseExperimentalCommandParser, envs.ZenlogUseExperimentalCommandParser)
	overwriteBoolWithEnviron(&c.UseSplice, envs.ZenlogUseSplice)

	if c.Maxproc < 1 {
		c.Maxproc = 1
	}

	if c.StartCommand == "" {
		shell := os.Getenv("SHELL")
		if shell != "" {
			c.StartCommand = "exec " + shell + " -l"
		} else {
			c.StartCommand = "exec /bin/sh -l"
		}
	}

	if c.TempDir == "" || !fileutils.DirExists(c.TempDir) {
		c.TempDir = utils.FirstNonEmpty(os.Getenv("TEMP"), os.Getenv("TMP"), "/tmp/")
	}

	ensureSlash(&c.LogDir)
	ensureSlash(&c.TempDir)

	// For E2E testing, override the PID with _ZENLOG_LOGGER_PID, if set.
	c.ZenlogPid = util.GetIntEnv("_ZENLOG_LOGGER_PID", os.Getpid())

	runtime.GOMAXPROCS(c.Maxproc)

	util.Dump("Logger config=", c)

	return &c
}

// InitConfigForCommands returns a Config for subcommands, loading from ~/.zenlog.toml and the environmental variables.
// Some of the parameters (such as ZenlogPid) will be inherited from the current zenlog session.
func InitConfigForCommands() *Config {
	var c Config

	// Take over some of the parameters from the logger.

	pid, err := strconv.Atoi(os.Getenv(envs.ZenlogPid))
	util.Check(err, "ZENLOG_PID not integer")
	c.ZenlogPid = pid

	c.LogDir = os.Getenv(envs.ZenlogDir)
	c.SourceDir = os.Getenv(envs.ZenlogSourceDir)
	c.OuterTty = os.Getenv(envs.ZenlogOuterTty)
	c.LoggerIn = os.Getenv(envs.ZenlogLoggerIn)
	c.LoggerOut = os.Getenv(envs.ZenlogLoggerOut)

	if c.ZenlogPid == 0 {
		util.Fatalf(envs.ZenlogPid + " not set.")
	}
	if c.LogDir == "" {
		util.Fatalf(envs.ZenlogDir + " not set.")
	}
	if c.OuterTty == "" {
		util.Fatalf(envs.ZenlogOuterTty + " not set.")
	}
	if c.LoggerIn == "" {
		util.Fatalf(envs.ZenlogLoggerIn + " not set.")
	}
	if c.LoggerOut == "" {
		util.Fatalf(envs.ZenlogLoggerOut + " not set.")
	}

	// We still need to load certain parameters from TOML.
	lc := InitConfigForLogger()
	c.AlwaysNoLogCommands = lc.AlwaysNoLogCommands
	c.PrefixCommands = lc.PrefixCommands
	c.CommandSplitter = lc.CommandSplitter
	c.CommentSplitter = lc.CommentSplitter
	c.UseExperimentalCommandParser = lc.UseExperimentalCommandParser
	c.TempDir = lc.TempDir
	c.Maxproc = lc.Maxproc

	util.Dump("Command config=", c)
	return &c
}

// GetConfig returns the singleton config suitable for the current run mode.
func GetConfig() *Config {
	if config == nil {
		if isLogger {
			config = InitConfigForLogger()
		} else {
			config = InitConfigForCommands()
		}
	}
	return config
}

// ZenlogSrcTopDir returns the fullpath of the source top directory.
func ZenlogSrcTopDir() string {
	zenlogBinDir := util.FindZenlogBinDir()

	configSourceDir := GetConfig().SourceDir
	if fileutils.DirExists(configSourceDir + "/subcommands") {
		return configSourceDir
	}

	for _, d := range utils.StringSlice("/../", "/../zenlog/", "/../src/github.com/omakoto/zenlog/") {
		candidate := zenlogBinDir + d
		candidate, err := filepath.Abs(candidate)
		util.Check(err, "Abs failed")

		if fileutils.DirExists(candidate + "/subcommands") {
			return candidate
		}
	}
	log.Fatalf("Zenlog source directory not found.")
	return ""
}
<commit_msg>Fix external subcommand launches outside of a zenlog session<commit_after>package config

import (
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/BurntSushi/toml"
	"github.com/omakoto/go-common/src/fileutils"
	"github.com/omakoto/go-common/src/utils"
	"github.com/omakoto/zenlog/zenlog/envs"
	"github.com/omakoto/zenlog/zenlog/util"
	"runtime"
)

var (
	isLogger    = false
	isLoggerSet = false
	config      *Config
)

func SetIsLogger(logger bool) {
	if isLoggerSet {
		util.Fatalf("SetIsLogger already set")
	}
	isLogger = logger
	isLoggerSet = true
}

// Config represents configuration parameters, read from ~/.zenlog.toml and overridden with the environmental variables.
type Config struct {
	LogDir              string `toml:"ZENLOG_DIR"`
	TempDir             string `toml:"ZENLOG_TEMP"`
	SourceDir           string `toml:"ZENLOG_SOURCE_DIR"`
	StartCommand        string `toml:"ZENLOG_START_COMMAND"`
	PrefixCommands      string `toml:"ZENLOG_PREFIX_COMMANDS"`
	AlwaysNoLogCommands string `toml:"ZENLOG_ALWAYS_NO_LOG_COMMANDS"`
	AutoFlush           bool   `toml:"ZENLOG_AUTO_FLUSH"`

	UseExperimentalCommandParser bool   `toml:"ZENLOG_USE_EXPERIMENTAL_COMMAND_PARSER"`
	CommandSplitter              string `toml:"ZENLOG_COMMAND_SPLITTER"`
	CommentSplitter              string `toml:"ZENLOG_COMMENT_SPLITTER"`

	Maxproc                 int `toml:"ZENLOG_MAXPROC"`
	CriticalCrashMaxSeconds int `toml:"ZENLOG_CRITICAL_CRASH_MAX_SECONDS"`

	UseSplice bool `toml:"ZENLOG_USE_SPLICE"`

	ZenlogPid int

	OuterTty  string
	LoggerIn  string
	LoggerOut string
}

// If a string is empty, get from a given environmental variabne.
func overwriteWithEnviron(to *string, envKey string, def string) {
	if val := os.Getenv(envKey); val != "" {
		*to = val
	} else {
		*to = os.ExpandEnv(*to)
	}
	if *to == "" {
		*to = def
	}
}

func overwriteBoolWithEnviron(to *bool, envKey string) {
	if val := os.Getenv(envKey); val != "" {
		*to = (val == "1")
	}
}

// Make sure a string ends with a slash.
func ensureSlash(v *string) {
	if strings.HasSuffix(*v, "/") {
		return
	}
	*v = *v + "/"
}

// InitConfigForLogger returns a Config for a new session, loading from ~/.zenlog.toml and the environmental variables.
func InitConfigForLogger() *Config {
	file := os.Getenv(envs.ZenlogConf)
	if file == "" {
		file = os.ExpandEnv("$HOME/.zenlog.toml")
	}
	var c Config

	c.UseExperimentalCommandParser = true // Default to true.
	c.UseSplice = true

	util.Debugf("config=%s", file)

	data, err := ioutil.ReadFile(file)
	if err == nil {
		if _, err := toml.Decode(string(data), &c); err != nil {
			util.Fatalf("Unable to load %s: %s", file, err)
		}
	} else if os.IsNotExist(err) {
		util.Warn(err, "%s doesn't exist; using the default instead", file)
	}

	overwriteWithEnviron(&c.StartCommand, "ZENLOG_START_COMMAND", "")
	overwriteWithEnviron(&c.LogDir, envs.ZenlogDir, os.ExpandEnv("$HOME/zenlog/"))
	overwriteWithEnviron(&c.SourceDir, envs.ZenlogSourceDir, "")
	overwriteWithEnviron(&c.PrefixCommands, "ZENLOG_PREFIX_COMMANDS", `(?:command|builtin|time|sudo|[a-zA-Z0-9_]+\=.*)`)
	overwriteWithEnviron(&c.AlwaysNoLogCommands, "ZENLOG_ALWAYS_NO_LOG_COMMANDS", `(?:vi|vim|man|nano|pico|emacs|zenlog.*)`)

	overwriteWithEnviron(&c.CommandSplitter, "ZENLOG_COMMAND_SPLITTER", "")
	overwriteWithEnviron(&c.CommentSplitter, "ZENLOG_COMMENT_SPLITTER", "")

	overwriteWithEnviron(&c.TempDir, envs.ZenlogTemp, "")

	overwriteBoolWithEnviron(&c.AutoFlush, envs.ZenlogAutoFlush)
	overwriteBoolWithEnviron(&c.UseExperimentalCommandParser, envs.ZenlogUseExperimentalCommandParser)
	overwriteBoolWithEnviron(&c.UseSplice, envs.ZenlogUseSplice)

	if c.Maxproc < 1 {
		c.Maxproc = 1
	}

	if c.StartCommand == "" {
		shell := os.Getenv("SHELL")
		if shell != "" {
			c.StartCommand = "exec " + shell + " -l"
		} else {
			c.StartCommand = "exec /bin/sh -l"
		}
	}

	if c.TempDir == "" || !fileutils.DirExists(c.TempDir) {
		c.TempDir = utils.FirstNonEmpty(os.Getenv("TEMP"), os.Getenv("TMP"), "/tmp/")
	}

	ensureSlash(&c.LogDir)
	ensureSlash(&c.TempDir)

	// For E2E testing, override the PID with _ZENLOG_LOGGER_PID, if set.
	c.ZenlogPid = util.GetIntEnv("_ZENLOG_LOGGER_PID", os.Getpid())

	runtime.GOMAXPROCS(c.Maxproc)

	util.Dump("Logger config=", c)

	return &c
}

// InitConfigForCommands returns a Config for subcommands, loading from ~/.zenlog.toml and the environmental variables.
// Some of the parameters (such as ZenlogPid) will be inherited from the current zenlog session.
func InitConfigForCommands() *Config {
	var c Config

	// Take over some of the parameters from the logger.

	pid, err := strconv.Atoi(os.Getenv(envs.ZenlogPid))
	util.Check(err, "ZENLOG_PID not integer")
	c.ZenlogPid = pid

	c.LogDir = os.Getenv(envs.ZenlogDir)
	c.SourceDir = os.Getenv(envs.ZenlogSourceDir)
	c.OuterTty = os.Getenv(envs.ZenlogOuterTty)
	c.LoggerIn = os.Getenv(envs.ZenlogLoggerIn)
	c.LoggerOut = os.Getenv(envs.ZenlogLoggerOut)

	if c.ZenlogPid == 0 {
		util.Fatalf(envs.ZenlogPid + " not set.")
	}
	if c.LogDir == "" {
		util.Fatalf(envs.ZenlogDir + " not set.")
	}
	if c.OuterTty == "" {
		util.Fatalf(envs.ZenlogOuterTty + " not set.")
	}
	if c.LoggerIn == "" {
		util.Fatalf(envs.ZenlogLoggerIn + " not set.")
	}
	if c.LoggerOut == "" {
		util.Fatalf(envs.ZenlogLoggerOut + " not set.")
	}

	// We still need to load certain parameters from TOML.
	lc := InitConfigForLogger()
	c.AlwaysNoLogCommands = lc.AlwaysNoLogCommands
	c.PrefixCommands = lc.PrefixCommands
	c.CommandSplitter = lc.CommandSplitter
	c.CommentSplitter = lc.CommentSplitter
	c.UseExperimentalCommandParser = lc.UseExperimentalCommandParser
	c.TempDir = lc.TempDir
	c.Maxproc = lc.Maxproc

	util.Dump("Command config=", c)
	return &c
}

// GetConfig returns the singleton config suitable for the current run mode.
func GetConfig() *Config {
	if config == nil {
		if isLogger {
			config = InitConfigForLogger()
		} else {
			config = InitConfigForCommands()
		}
	}
	return config
}

// ZenlogSrcTopDir returns the fullpath of the source top directory.
func ZenlogSrcTopDir() string {
	zenlogBinDir := util.FindZenlogBinDir()

	// Note, this needs to work even outside of a zenlog session, where InitConfigForCommands()
	// doesn't work, so we need to use InitConfigForLogger().
	config = InitConfigForLogger()
	configSourceDir := config.SourceDir
	if fileutils.DirExists(configSourceDir + "/subcommands") {
		return configSourceDir
	}

	for _, d := range utils.StringSlice("/../", "/../zenlog/", "/../src/github.com/omakoto/zenlog/") {
		candidate := zenlogBinDir + d
		candidate, err := filepath.Abs(candidate)
		util.Check(err, "Abs failed")

		if fileutils.DirExists(candidate + "/subcommands") {
			return candidate
		}
	}
	log.Fatalf("Zenlog source directory not found.")
	return ""
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Ashley Jeffs
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package reader

import (
	"strings"
	"sync"
	"time"

	"github.com/Jeffail/benthos/lib/log"
	"github.com/Jeffail/benthos/lib/message"
	"github.com/Jeffail/benthos/lib/metrics"
	"github.com/Jeffail/benthos/lib/types"
	nats "github.com/nats-io/go-nats"
)

//------------------------------------------------------------------------------

// NATSConfig contains configuration fields for the NATS input type.
type NATSConfig struct {
	URLs    []string `json:"urls" yaml:"urls"`
	Subject string   `json:"subject" yaml:"subject"`
}

// NewNATSConfig creates a new NATSConfig with default values.
func NewNATSConfig() NATSConfig {
	return NATSConfig{
		URLs:    []string{nats.DefaultURL},
		Subject: "benthos_messages",
	}
}

//------------------------------------------------------------------------------

// NATS is an input type that receives NATS messages.
type NATS struct {
	urls  string
	conf  NATSConfig
	stats metrics.Type
	log   log.Modular

	cMut sync.Mutex

	natsConn      *nats.Conn
	natsSub       *nats.Subscription
	natsChan      chan *nats.Msg
	interruptChan chan struct{}
}

// NewNATS creates a new NATS input type.
func NewNATS(conf NATSConfig, log log.Modular, stats metrics.Type) (Type, error) {
	n := NATS{
		conf:          conf,
		stats:         stats,
		log:           log.NewModule(".input.nats"),
		interruptChan: make(chan struct{}),
	}
	n.urls = strings.Join(conf.URLs, ",")

	return &n, nil
}

//------------------------------------------------------------------------------

// Connect establishes a connection to a NATS server.
func (n *NATS) Connect() error {
	n.cMut.Lock()
	defer n.cMut.Unlock()

	if n.natsConn != nil {
		return nil
	}

	var natsConn *nats.Conn
	var natsSub *nats.Subscription
	var err error

	if natsConn, err = nats.Connect(n.urls); err != nil {
		return err
	}
	natsChan := make(chan *nats.Msg)
	if natsSub, err = natsConn.ChanSubscribe(n.conf.Subject, natsChan); err != nil {
		return err
	}

	n.log.Infof("Receiving NATS messages from URLs: %s\n", n.urls)

	n.natsConn = natsConn
	n.natsSub = natsSub
	n.natsChan = natsChan
	return nil
}

func (n *NATS) disconnect() {
	n.cMut.Lock()
	defer n.cMut.Unlock()

	if n.natsSub != nil {
		n.natsSub.Unsubscribe()
		n.natsSub = nil
	}
	if n.natsConn != nil {
		n.natsConn.Close()
		n.natsConn = nil
	}
	n.natsChan = nil
}

// Read attempts to read a new message from the NATS subject.
func (n *NATS) Read() (types.Message, error) {
	n.cMut.Lock()
	natsChan := n.natsChan
	n.cMut.Unlock()

	var msg *nats.Msg
	var open bool
	select {
	case msg, open = <-natsChan:
	case _, open = <-n.interruptChan:
	}
	if !open {
		n.disconnect()
		return nil, types.ErrNotConnected
	}

	bmsg := message.New([][]byte{msg.Data})
	bmsg.Get(0).Metadata().Set("nats_subject", msg.Subject)

	return bmsg, nil
}

// Acknowledge is a noop since NATS messages do not support acknowledgments.
func (n *NATS) Acknowledge(err error) error {
	return nil
}

// CloseAsync shuts down the NATS input and stops processing requests.
func (n *NATS) CloseAsync() {
	close(n.interruptChan)
}

// WaitForClose blocks until the NATS input has closed down.
func (n *NATS) WaitForClose(timeout time.Duration) error {
	return nil
}

//------------------------------------------------------------------------------
<commit_msg>adding nats queue as an option<commit_after>// Copyright (c) 2018 Ashley Jeffs
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package reader

import (
	"strings"
	"sync"
	"time"

	"github.com/Jeffail/benthos/lib/log"
	"github.com/Jeffail/benthos/lib/message"
	"github.com/Jeffail/benthos/lib/metrics"
	"github.com/Jeffail/benthos/lib/types"
	nats "github.com/nats-io/go-nats"
)

//------------------------------------------------------------------------------

// NATSConfig contains configuration fields for the NATS input type.
type NATSConfig struct {
	URLs    []string `json:"urls" yaml:"urls"`
	Subject string   `json:"subject" yaml:"subject"`
	QueueID string   `json:"queue" yaml:"queue"`
}

// NewNATSConfig creates a new NATSConfig with default values.
func NewNATSConfig() NATSConfig {
	return NATSConfig{
		URLs:    []string{nats.DefaultURL},
		Subject: "benthos_messages",
		QueueID: "benthos_queue",
	}
}

//------------------------------------------------------------------------------

// NATS is an input type that receives NATS messages.
type NATS struct {
	urls  string
	conf  NATSConfig
	stats metrics.Type
	log   log.Modular

	cMut sync.Mutex

	natsConn      *nats.Conn
	natsSub       *nats.Subscription
	natsChan      chan *nats.Msg
	interruptChan chan struct{}
}

// NewNATS creates a new NATS input type.
func NewNATS(conf NATSConfig, log log.Modular, stats metrics.Type) (Type, error) {
	n := NATS{
		conf:          conf,
		stats:         stats,
		log:           log.NewModule(".input.nats"),
		interruptChan: make(chan struct{}),
	}
	n.urls = strings.Join(conf.URLs, ",")

	return &n, nil
}

//------------------------------------------------------------------------------

// Connect establishes a connection to a NATS server.
func (n *NATS) Connect() error {
	n.cMut.Lock()
	defer n.cMut.Unlock()

	if n.natsConn != nil {
		return nil
	}

	var natsConn *nats.Conn
	var natsSub *nats.Subscription
	var err error

	if natsConn, err = nats.Connect(n.urls); err != nil {
		return err
	}
	natsChan := make(chan *nats.Msg)

	if len(n.conf.QueueID) > 0 {
		natsSub, err = natsConn.ChanQueueSubscribe(n.conf.Subject, n.conf.QueueID, natsChan)
	} else {
		natsSub, err = natsConn.ChanSubscribe(n.conf.Subject, natsChan)
	}

	if err != nil {
		return err
	}

	n.log.Infof("Receiving NATS messages from URLs: %s\n", n.urls)

	n.natsConn = natsConn
	n.natsSub = natsSub
	n.natsChan = natsChan
	return nil
}

func (n *NATS) disconnect() {
	n.cMut.Lock()
	defer n.cMut.Unlock()

	if n.natsSub != nil {
		n.natsSub.Unsubscribe()
		n.natsSub = nil
	}
	if n.natsConn != nil {
		n.natsConn.Close()
		n.natsConn = nil
	}
	n.natsChan = nil
}

// Read attempts to read a new message from the NATS subject.
func (n *NATS) Read() (types.Message, error) {
	n.cMut.Lock()
	natsChan := n.natsChan
	n.cMut.Unlock()

	var msg *nats.Msg
	var open bool
	select {
	case msg, open = <-natsChan:
	case _, open = <-n.interruptChan:
	}
	if !open {
		n.disconnect()
		return nil, types.ErrNotConnected
	}

	bmsg := message.New([][]byte{msg.Data})
	bmsg.Get(0).Metadata().Set("nats_subject", msg.Subject)

	return bmsg, nil
}

// Acknowledge is a noop since NATS messages do not support acknowledgments.
func (n *NATS) Acknowledge(err error) error {
	return nil
}

// CloseAsync shuts down the NATS input and stops processing requests.
func (n *NATS) CloseAsync() {
	close(n.interruptChan)
}

// WaitForClose blocks until the NATS input has closed down.
func (n *NATS) WaitForClose(timeout time.Duration) error {
	return nil
}

//------------------------------------------------------------------------------
<|endoftext|>


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

package main

import (
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// TODO: Check CPU performance governor before each benchmark.

// TODO: Mode to prioritize commits around big performance changes.

// TODO: Support running pre-built binaries without specific hashes.
// This is useful for testing things that aren't yet committed or that
// require unusual build steps.

// TODO: If we switched to the extended benchmark format and writing
// out one big file, we could count runs using the configuration
// blocks instead of the silly "PASS" search we use now.

var run struct {
	order      string
	topLevel   string
	benchFlags string
	buildCmd   string
	iterations int
	saveTree   bool
}

var cmdRunFlags = flag.NewFlagSet(os.Args[0]+" run", flag.ExitOnError)

func init() {
	isXBenchmark := false
	if abs, _ := os.Getwd(); strings.HasSuffix(abs, "golang.org/x/benchmarks/bench") {
		isXBenchmark = true
	}

	f := cmdRunFlags
	f.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s run [flags] <revision range>\n", os.Args[0])
		f.PrintDefaults()
	}
	f.StringVar(&run.order, "order", "seq", "run benchmarks in `order`, which must be one of: seq, spread")
	f.StringVar(&gitDir, "C", "", "run git in `dir`")
	defaultBenchFlags := "-test.run NONE -test.bench ."
	if isXBenchmark {
		defaultBenchFlags = ""
	}
	f.StringVar(&run.benchFlags, "benchflags", defaultBenchFlags, "pass `flags` to benchmark")
	defaultBuildCmd := "go test -c"
	if isXBenchmark {
		defaultBuildCmd = "go build"
	}
	f.StringVar(&run.buildCmd, "buildcmd", defaultBuildCmd, "build benchmark using \"`cmd` -o <bin>\"")
	f.IntVar(&run.iterations, "n", 5, "run each benchmark `N` times")
	f.StringVar(&outDir, "o", "", "write binaries and logs to `directory`")
	f.BoolVar(&run.saveTree, "save-tree", false, "save Go trees using gover and run benchmarks under saved trees")
	f.BoolVar(&dryRun, "dry-run", false, "print commands but do not run them")
	registerSubcommand("run", "[flags] <revision range> - run benchmarks", cmdRun, f)
}

func cmdRun() {
	if cmdRunFlags.NArg() < 1 {
		cmdRunFlags.Usage()
		os.Exit(2)
	}

	var pickCommit func([]*commitInfo) *commitInfo
	switch run.order {
	case "seq":
		pickCommit = pickCommitSeq
	case "spread":
		pickCommit = pickCommitSpread
	default:
		fmt.Fprintf(os.Stderr, "unknown order: %s\n", run.order)
		cmdRunFlags.Usage()
		os.Exit(2)
	}

	commits := getCommits(cmdRunFlags.Args())

	// Get other git information.
	run.topLevel = trimNL(git("rev-parse", "--show-toplevel"))

	status := NewStatusReporter()
	defer status.Stop()

	for {
		doneIters, totalIters, partialCommits, doneCommits, failedCommits := runStats(commits)
		unstartedCommits := len(commits) - (partialCommits + doneCommits + failedCommits)
		msg := fmt.Sprintf("%d/%d runs, %d unstarted+%d partial+%d done+%d failed commits", doneIters, totalIters, unstartedCommits, partialCommits, doneCommits, failedCommits)
		// TODO: Count builds and runs separately.
		status.Progress(msg, float64(doneIters)/float64(totalIters))

		commit := pickCommit(commits)
		if commit == nil {
			break
		}
		runBenchmark(commit, status)
	}
}

func runStats(commits []*commitInfo) (doneIters, totalIters, partialCommits, doneCommits, failedCommits int) {
	for _, c := range commits {
		if c.count >= run.iterations {
			// Don't care if it failed.
			doneIters += c.count
			totalIters += c.count
		} else if c.runnable() {
			doneIters += c.count
			totalIters += run.iterations
		}

		if c.count == run.iterations {
			doneCommits++
		} else if c.runnable() {
			if c.count != 0 {
				partialCommits++
			}
		} else {
			failedCommits++
		}
	}
	return
}

// pickCommitSeq picks the next commit to run based on the most recent
// commit with the fewest iterations.
func pickCommitSeq(commits []*commitInfo) *commitInfo {
	var minCommit *commitInfo
	for _, commit := range commits {
		if !commit.runnable() {
			continue
		}
		if minCommit == nil || commit.count < minCommit.count {
			minCommit = commit
		}
	}
	return minCommit
}

// pickCommitSpread picks the next commit to run from commits using an
// algorithm that spreads out the runs.
func pickCommitSpread(commits []*commitInfo) *commitInfo {
	// Assign weights to each commit. This is thoroughly
	// heuristic, but it's geared toward either increasing the
	// iteration count of commits that we have, or picking a new
	// commit so as to spread out the commits we have.
	weights := make([]int, len(commits))
	totalWeight := 0

	nPartial := 0
	for _, commit := range commits {
		if commit.partial() {
			nPartial++
		}
	}
	if nPartial >= len(commits)/10 {
		// Limit the number of partially completed revisions
		// to 10% by only choosing a partial commit in this
		// case.
		for i, commit := range commits {
			if commit.partial() {
				// Bias toward commits that are
				// further from done.
				weights[i] = run.iterations - commit.count
			}
		}
	} else {
		// Pick a new commit weighted by its distance from a
		// commit that we already have.

		// Find distance from left to right.
		distance := len(commits)
		haveAny := false
		for i, commit := range commits {
			if commit.count > 0 {
				distance = 1
				haveAny = true
			} else if commit.runnable() {
				distance++
			}
			weights[i] = distance
		}

		// Find distance from right to left.
		distance = len(commits)
		for i := len(commits) - 1; i >= 0; i-- {
			commit := commits[i]
			if commit.count > 0 {
				distance = 1
			} else if commit.runnable() {
				distance++
			}

			if distance < weights[i] {
				weights[i] = distance
			}
		}

		if !haveAny {
			// We don't have any commits. Pick one uniformly.
			for i := range commits {
				weights[i] = 1
			}
		}

		// Zero non-runnable commits.
		for i, commit := range commits {
			if !commit.runnable() {
				weights[i] = 0
			}
		}
	}

	for _, w := range weights {
		totalWeight += w
	}
	if totalWeight == 0 {
		return nil
	}

	// Pick a commit based on the weights.
	x := rand.Intn(totalWeight)
	cumulative := 0
	for i, w := range weights {
		cumulative += w
		if cumulative > x {
			return commits[i]
		}
	}
	panic("unreachable")
}

// runBenchmark runs the benchmark at commit. It updates commit.count,
// commit.fails, and commit.buildFailed as appropriate and writes to
// the commit log to record the outcome.
func runBenchmark(commit *commitInfo, status *StatusReporter) {
	// Build the benchmark if necessary.
	if !exists(commit.binPath) {
		runStatus(status, commit, "building")

		// Check out the appropriate commit. This is necessary
		// even if we're using gover because the benchmark
		// itself might have changed (e.g., bug fixes).
		git("checkout", "-q", commit.hash)

		var buildCmd []string
		if commit.gover {
			// TODO: It would be better if gover took a
			// long hash and did the right thing.
			buildCmd = []string{"gover", "run", commit.hash[:7]}
		} else {
			// If this is the Go toolchain, do a full
			// make.bash. Otherwise, we assume that go
			// test -c will build the necessary
			// dependencies.
			if exists(filepath.Join(run.topLevel, "src", "make.bash")) {
				cmd := exec.Command("./make.bash")
				cmd.Dir = filepath.Join(run.topLevel, "src")
				if dryRun {
					dryPrint(cmd)
				} else if out, err := cmd.CombinedOutput(); err != nil {
					detail := indent(string(out)) + indent(err.Error())
					fmt.Fprintf(os.Stderr, "failed to build toolchain at %s:\n%s", commit.hash, detail)
					commit.buildFailed = true
					commit.writeLog([]byte("BUILD FAILED:\n" + detail))
					return
				}
				if run.saveTree && doGoverSave() == nil {
					commit.gover = true
				}
			}
			// Assume build command is in $PATH.
			//
			// TODO: Force PATH if we built the toolchain.
			buildCmd = []string{}
		}

		buildCmd = append(buildCmd, strings.Fields(run.buildCmd)...)
		buildCmd = append(buildCmd, "-o", commit.binPath)
		cmd := exec.Command(buildCmd[0], buildCmd[1:]...)
		if dryRun {
			dryPrint(cmd)
		} else if out, err := cmd.CombinedOutput(); err != nil {
			detail := indent(string(out)) + indent(err.Error())
			fmt.Fprintf(os.Stderr, "failed to build tests at %s:\n%s", commit.hash, detail)
			commit.buildFailed = true
			commit.writeLog([]byte("BUILD FAILED:\n" + detail))
			return
		}
	}

	// Run the benchmark.
	runStatus(status, commit, "running")
	name := commit.binPath
	if filepath.Base(name) == name {
		// Make exec.Command treat this as a relative path.
		name = "./" + name
	}
	args := append([]string{name}, strings.Fields(run.benchFlags)...)
	if run.saveTree {
		args = append([]string{"gover", "run", commit.hash[:7]}, args...)
	}
	cmd := exec.Command(args[0], args[1:]...)
	if dryRun {
		dryPrint(cmd)
		commit.count++
		return
	}
	out, err := cmd.CombinedOutput()
	if err == nil {
		commit.count++
		if c, f, _ := countRuns(bytes.NewBuffer(out)); c+f == 0 {
			// This log doesn't count as a run, probably
			// because it's missing a "PASS". Add one so
			// we can read this back in properly later.
			if !bytes.HasSuffix(out, []byte{'\n'}) {
				out = append(out, '\n')
			}
			out = append(out, []byte("PASS\n")...)
		}
	} else {
		fmt.Fprintf(os.Stderr, "failed to run benchmark at %s:\n%s", commit.hash, out)
		commit.fails++
		// Indent the output so we don't get confused by
		// benchmarks in it or "PASS" lines.
		out = []byte("FAILED:\n" + indent(string(out)) + indent(err.Error()))
	}

	// Write the benchmark output.
	commit.writeLog(out)
}

func doGoverSave() error {
	cmd := exec.Command("gover", "save")
	cmd.Dir = run.topLevel
	if dryRun {
		dryPrint(cmd)
		return nil
	} else {
		out, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Fprintf(os.Stderr, "gover save failed: %s:\n%s", err, indent(string(out)))
		}
		return err
	}
}

// runStatus updates the status message for commit.
func runStatus(sr *StatusReporter, commit *commitInfo, status string) {
	sr.Message(fmt.Sprintf("commit %s, iteration %d/%d: %s...", commit.hash[:7], commit.count+1, run.iterations, status))
}
<commit_msg>benchmany: enforce 30m timeout on benchmarks<commit_after>// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

// TODO: Check CPU performance governor before each benchmark.

// TODO: Mode to prioritize commits around big performance changes.

// TODO: Support running pre-built binaries without specific hashes.
// This is useful for testing things that aren't yet committed or that
// require unusual build steps.

// TODO: If we switched to the extended benchmark format and writing
// out one big file, we could count runs using the configuration
// blocks instead of the silly "PASS" search we use now.

var run struct {
	order      string
	topLevel   string
	benchFlags string
	buildCmd   string
	iterations int
	saveTree   bool
	timeout    time.Duration
}

var cmdRunFlags = flag.NewFlagSet(os.Args[0]+" run", flag.ExitOnError)

func init() {
	isXBenchmark := false
	if abs, _ := os.Getwd(); strings.HasSuffix(abs, "golang.org/x/benchmarks/bench") {
		isXBenchmark = true
	}

	f := cmdRunFlags
	f.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s run [flags] <revision range>\n", os.Args[0])
		f.PrintDefaults()
	}
	f.StringVar(&run.order, "order", "seq", "run benchmarks in `order`, which must be one of: seq, spread")
	f.StringVar(&gitDir, "C", "", "run git in `dir`")
	defaultBenchFlags := "-test.run NONE -test.bench ."
	if isXBenchmark {
		defaultBenchFlags = ""
	}
	f.StringVar(&run.benchFlags, "benchflags", defaultBenchFlags, "pass `flags` to benchmark")
	defaultBuildCmd := "go test -c"
	if isXBenchmark {
		defaultBuildCmd = "go build"
	}
	f.StringVar(&run.buildCmd, "buildcmd", defaultBuildCmd, "build benchmark using \"`cmd` -o <bin>\"")
	f.IntVar(&run.iterations, "n", 5, "run each benchmark `N` times")
	f.StringVar(&outDir, "o", "", "write binaries and logs to `directory`")
	f.BoolVar(&run.saveTree, "save-tree", false, "save Go trees using gover and run benchmarks under saved trees")
	f.DurationVar(&run.timeout, "timeout", 30*time.Minute, "time out a run after `duration`")
	f.BoolVar(&dryRun, "dry-run", false, "print commands but do not run them")
	registerSubcommand("run", "[flags] <revision range> - run benchmarks", cmdRun, f)
}

func cmdRun() {
	if cmdRunFlags.NArg() < 1 {
		cmdRunFlags.Usage()
		os.Exit(2)
	}

	var pickCommit func([]*commitInfo) *commitInfo
	switch run.order {
	case "seq":
		pickCommit = pickCommitSeq
	case "spread":
		pickCommit = pickCommitSpread
	default:
		fmt.Fprintf(os.Stderr, "unknown order: %s\n", run.order)
		cmdRunFlags.Usage()
		os.Exit(2)
	}

	commits := getCommits(cmdRunFlags.Args())

	// Get other git information.
	run.topLevel = trimNL(git("rev-parse", "--show-toplevel"))

	status := NewStatusReporter()
	defer status.Stop()

	for {
		doneIters, totalIters, partialCommits, doneCommits, failedCommits := runStats(commits)
		unstartedCommits := len(commits) - (partialCommits + doneCommits + failedCommits)
		msg := fmt.Sprintf("%d/%d runs, %d unstarted+%d partial+%d done+%d failed commits", doneIters, totalIters, unstartedCommits, partialCommits, doneCommits, failedCommits)
		// TODO: Count builds and runs separately.
		status.Progress(msg, float64(doneIters)/float64(totalIters))

		commit := pickCommit(commits)
		if commit == nil {
			break
		}
		runBenchmark(commit, status)
	}
}

func runStats(commits []*commitInfo) (doneIters, totalIters, partialCommits, doneCommits, failedCommits int) {
	for _, c := range commits {
		if c.count >= run.iterations {
			// Don't care if it failed.
			doneIters += c.count
			totalIters += c.count
		} else if c.runnable() {
			doneIters += c.count
			totalIters += run.iterations
		}

		if c.count == run.iterations {
			doneCommits++
		} else if c.runnable() {
			if c.count != 0 {
				partialCommits++
			}
		} else {
			failedCommits++
		}
	}
	return
}

// pickCommitSeq picks the next commit to run based on the most recent
// commit with the fewest iterations.
func pickCommitSeq(commits []*commitInfo) *commitInfo {
	var minCommit *commitInfo
	for _, commit := range commits {
		if !commit.runnable() {
			continue
		}
		if minCommit == nil || commit.count < minCommit.count {
			minCommit = commit
		}
	}
	return minCommit
}

// pickCommitSpread picks the next commit to run from commits using an
// algorithm that spreads out the runs.
func pickCommitSpread(commits []*commitInfo) *commitInfo {
	// Assign weights to each commit. This is thoroughly
	// heuristic, but it's geared toward either increasing the
	// iteration count of commits that we have, or picking a new
	// commit so as to spread out the commits we have.
	weights := make([]int, len(commits))
	totalWeight := 0

	nPartial := 0
	for _, commit := range commits {
		if commit.partial() {
			nPartial++
		}
	}
	if nPartial >= len(commits)/10 {
		// Limit the number of partially completed revisions
		// to 10% by only choosing a partial commit in this
		// case.
		for i, commit := range commits {
			if commit.partial() {
				// Bias toward commits that are
				// further from done.
				weights[i] = run.iterations - commit.count
			}
		}
	} else {
		// Pick a new commit weighted by its distance from a
		// commit that we already have.

		// Find distance from left to right.
		distance := len(commits)
		haveAny := false
		for i, commit := range commits {
			if commit.count > 0 {
				distance = 1
				haveAny = true
			} else if commit.runnable() {
				distance++
			}
			weights[i] = distance
		}

		// Find distance from right to left.
		distance = len(commits)
		for i := len(commits) - 1; i >= 0; i-- {
			commit := commits[i]
			if commit.count > 0 {
				distance = 1
			} else if commit.runnable() {
				distance++
			}

			if distance < weights[i] {
				weights[i] = distance
			}
		}

		if !haveAny {
			// We don't have any commits. Pick one uniformly.
			for i := range commits {
				weights[i] = 1
			}
		}

		// Zero non-runnable commits.
		for i, commit := range commits {
			if !commit.runnable() {
				weights[i] = 0
			}
		}
	}

	for _, w := range weights {
		totalWeight += w
	}
	if totalWeight == 0 {
		return nil
	}

	// Pick a commit based on the weights.
	x := rand.Intn(totalWeight)
	cumulative := 0
	for i, w := range weights {
		cumulative += w
		if cumulative > x {
			return commits[i]
		}
	}
	panic("unreachable")
}

// runBenchmark runs the benchmark at commit. It updates commit.count,
// commit.fails, and commit.buildFailed as appropriate and writes to
// the commit log to record the outcome.
func runBenchmark(commit *commitInfo, status *StatusReporter) {
	// Build the benchmark if necessary.
	if !exists(commit.binPath) {
		runStatus(status, commit, "building")

		// Check out the appropriate commit. This is necessary
		// even if we're using gover because the benchmark
		// itself might have changed (e.g., bug fixes).
		git("checkout", "-q", commit.hash)

		var buildCmd []string
		if commit.gover {
			// TODO: It would be better if gover took a
			// long hash and did the right thing.
			buildCmd = []string{"gover", "run", commit.hash[:7]}
		} else {
			// If this is the Go toolchain, do a full
			// make.bash. Otherwise, we assume that go
			// test -c will build the necessary
			// dependencies.
			if exists(filepath.Join(run.topLevel, "src", "make.bash")) {
				cmd := exec.Command("./make.bash")
				cmd.Dir = filepath.Join(run.topLevel, "src")
				if dryRun {
					dryPrint(cmd)
				} else if out, err := combinedOutputTimeout(cmd); err != nil {
					detail := indent(string(out)) + indent(err.Error())
					fmt.Fprintf(os.Stderr, "failed to build toolchain at %s:\n%s", commit.hash, detail)
					commit.buildFailed = true
					commit.writeLog([]byte("BUILD FAILED:\n" + detail))
					return
				}
				if run.saveTree && doGoverSave() == nil {
					commit.gover = true
				}
			}
			// Assume build command is in $PATH.
			//
			// TODO: Force PATH if we built the toolchain.
			buildCmd = []string{}
		}

		buildCmd = append(buildCmd, strings.Fields(run.buildCmd)...)
		buildCmd = append(buildCmd, "-o", commit.binPath)
		cmd := exec.Command(buildCmd[0], buildCmd[1:]...)
		if dryRun {
			dryPrint(cmd)
		} else if out, err := combinedOutputTimeout(cmd); err != nil {
			detail := indent(string(out)) + indent(err.Error())
			fmt.Fprintf(os.Stderr, "failed to build tests at %s:\n%s", commit.hash, detail)
			commit.buildFailed = true
			commit.writeLog([]byte("BUILD FAILED:\n" + detail))
			return
		}
	}

	// Run the benchmark.
	runStatus(status, commit, "running")
	name := commit.binPath
	if filepath.Base(name) == name {
		// Make exec.Command treat this as a relative path.
		name = "./" + name
	}
	args := append([]string{name}, strings.Fields(run.benchFlags)...)
	if run.saveTree {
		args = append([]string{"gover", "run", commit.hash[:7]}, args...)
	}
	cmd := exec.Command(args[0], args[1:]...)
	if dryRun {
		dryPrint(cmd)
		commit.count++
		return
	}
	out, err := combinedOutputTimeout(cmd)
	if err == nil {
		commit.count++
		if c, f, _ := countRuns(bytes.NewBuffer(out)); c+f == 0 {
			// This log doesn't count as a run, probably
			// because it's missing a "PASS". Add one so
			// we can read this back in properly later.
			if !bytes.HasSuffix(out, []byte{'\n'}) {
				out = append(out, '\n')
			}
			out = append(out, []byte("PASS\n")...)
		}
	} else {
		fmt.Fprintf(os.Stderr, "failed to run benchmark at %s:\n%s", commit.hash, out)
		commit.fails++
		// Indent the output so we don't get confused by
		// benchmarks in it or "PASS" lines.
		out = []byte("FAILED:\n" + indent(string(out)) + indent(err.Error()))
	}

	// Write the benchmark output.
	commit.writeLog(out)
}

func doGoverSave() error {
	cmd := exec.Command("gover", "save")
	cmd.Dir = run.topLevel
	if dryRun {
		dryPrint(cmd)
		return nil
	} else {
		out, err := cmd.CombinedOutput()
		if err != nil {
			fmt.Fprintf(os.Stderr, "gover save failed: %s:\n%s", err, indent(string(out)))
		}
		return err
	}
}

// runStatus updates the status message for commit.
func runStatus(sr *StatusReporter, commit *commitInfo, status string) {
	sr.Message(fmt.Sprintf("commit %s, iteration %d/%d: %s...", commit.hash[:7], commit.count+1, run.iterations, status))
}

// combinedOutputTimeout is like c.CombinedOutput(), but if
// run.timeout != 0, it will kill c after run.timeout time expires.
func combinedOutputTimeout(c *exec.Cmd) (out []byte, err error) {
	var b bytes.Buffer
	c.Stdout = &b
	c.Stderr = &b
	if err := c.Start(); err != nil {
		return nil, err
	}

	if run.timeout == 0 {
		err := c.Wait()
		return b.Bytes(), err
	}

	tick := time.NewTimer(run.timeout)
	trace := signalTrace
	done := make(chan error)
	go func() {
		done <- c.Wait()
	}()
loop:
	for {
		select {
		case err = <-done:
			break loop
		case <-tick.C:
			if trace != nil {
				fmt.Fprintf(os.Stderr, "command timed out; sending %v\n", trace)
				c.Process.Signal(trace)
				tick = time.NewTimer(5 * time.Second)
				trace = nil
			} else {
				fmt.Fprintf(os.Stderr, "command timed out; killing\n")
				c.Process.Kill()
			}
		}
	}
	tick.Stop()
	return b.Bytes(), err
}
<|endoftext|>


<commit_before>package web

import (
	"context"
	"crypto/tls"
	"errors"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	_app  *Application
	_once sync.Once
)

// Callback function
type Callback func(ctx *Context)

// PanicCallback function
type PanicCallback func(http.ResponseWriter, *http.Request, interface{})

// Application is type of a web.Application
type Application struct {
	trees       map[string]*node
	middlewares Middlewares
	logger      *log.Logger
	panic       PanicCallback
	paramsPool  sync.Pool
	maxParams   uint16
}

// Create return a singleton web.Application
func Create() *Application {
	_once.Do(func() {
		_app = newApplication()
	})
	return _app
}

// newApplication return a web.Application
func newApplication() *Application {
	app := &Application{
		middlewares: Middlewares{},
		logger:      log.New(os.Stdout, "", log.Ldate|log.Ltime),
	}
	return app
}

// SetLogger set Logger
func (app *Application) SetLogger(logger *log.Logger) {
	app.logger = logger
}

// SetPanic set Logger
func (app *Application) SetPanic(panic PanicCallback) {
	app.panic = panic
}

// Use Add the given callback function to this application.middlewares.
func (app *Application) Use(path string, callback Callback) {

	m := Middleware{
		Path:     cleanPath(path),
		Callback: callback,
	}

	app.middlewares = append(app.middlewares, m)
}

// Resource map controller path
func (app *Application) Resource(path string, controller Controller) {
	app.ResourceFn(path, controller, nil)
}

// ResourceFn map controller path and wrap fn
func (app *Application) ResourceFn(path string, controller Controller, fn func(cb Callback, keys ...string) Callback, keys ...string) {

	path = cleanPath(path)

	if fn == nil {
		app.Get(path, controller.Index)
		app.Post(path, controller.Create)
		app.Get(path+":id", controller.Detail)
		app.Patch(path+":id", controller.Update)
		app.Put(path+":id", controller.Update)
		app.Delete(path+":id", controller.Destroy)
	} else {
		app.Get(path, fn(controller.Index, keys...))
		app.Post(path, fn(controller.Create, keys...))
		app.Get(path+":id", fn(controller.Detail, keys...))
		app.Patch(path+":id", fn(controller.Update, keys...))
		app.Put(path+":id", fn(controller.Update, keys...))
		app.Delete(path+":id", fn(controller.Destroy, keys...))
	}
}

// On add event
func (app *Application) On(name string, cb Callback) {

}

// Get method
func (app *Application) Get(path string, cb Callback) {
	app.addRoute(http.MethodGet, path, cb)
}

// Post method
func (app *Application) Post(path string, cb Callback) {
	app.addRoute(http.MethodPost, path, cb)
}

// Put method
func (app *Application) Put(path string, cb Callback) {
	app.addRoute(http.MethodPut, path, cb)
}

// Patch method
func (app *Application) Patch(path string, cb Callback) {
	app.addRoute(http.MethodPatch, path, cb)
}

// Delete method
func (app *Application) Delete(path string, cb Callback) {
	app.addRoute(http.MethodDelete, path, cb)
}

// Options method
func (app *Application) Options(path string, cb Callback) {
	app.addRoute(http.MethodOptions, path, cb)
}

func (app *Application) addRoute(method, path string, cb Callback) {

	if method == "" {
		panic("method must not be empty")
	}

	if len(path) < 1 || path[0] != '/' {
		panic("path must begin with '/' in path '" + path + "'")
	}

	if cb == nil {
		panic("callback must not be nil")
	}

	if app.trees == nil {
		app.trees = make(map[string]*node)
	}

	root := app.trees[method]

	if root == nil {
		root = new(node)
		app.trees[method] = root
	}

	root.addRoute(path, cb)

	if pc := countParams(path); pc > app.maxParams {
		app.maxParams = pc
	}

	if app.paramsPool.New == nil && app.maxParams > 0 {
		app.paramsPool.New = func() interface{} {
			ps := make(Params, 0, app.maxParams)
			return &ps
		}
	}
}

// ServeHTTP
func (app *Application) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()
	defer app.recv(w, r)
	path := r.URL.Path

	if root := app.trees[r.Method]; root != nil {

		if callback, params, tsr := root.getValue(path, app.getParams); callback != nil {

			ctx := newContext(w, r, params)
			app.putParams(params)

			app.middlewares.exec(path, ctx)

			runTime := time.Now()
			callback(ctx)
			endTime := time.Now()

			app.logf("%s %s %s %s %s", r.Method, path, endTime.Sub(startTime), runTime.Sub(startTime), endTime.Sub(runTime))

			return

		} else if r.Method != http.MethodConnect && path != "/" {

			code := http.StatusMovedPermanently

			if r.Method != http.MethodGet {
				code = http.StatusPermanentRedirect
			}

			if tsr {

				if len(path) > 1 && path[len(path)-1] == '/' {
					r.URL.Path = path[:len(path)-1]
				} else {
					r.URL.Path = path + "/"
				}

				http.Redirect(w, r, r.URL.String(), code)

				return
			}
		}
	}

	http.NotFound(w, r)
}

// ListenAndServe Serve with options on addr
func (app *Application) ListenAndServe(addr string, cb func(*http.Server)) error {

	l, err := net.Listen("tcp", addr)

	if err != nil {
		log.Fatal("Listen:", err)
	}

	defer l.Close()

	return app.serve(addr, l, cb)
}

// ListenAndServeTLS Serve with tls and options on addr
func (app *Application) ListenAndServeTLS(addr string, tlsConfig *tls.Config, cb func(*http.Server)) error {

	l, err := tls.Listen("tcp", addr, tlsConfig)

	if err != nil {
		log.Fatal("Listen:", err)
	}

	defer l.Close()

	return app.serve(addr, l, cb)
}

func (app *Application) serve(addr string, listener net.Listener, cb func(*http.Server)) error {

	mux := http.NewServeMux()

	mux.Handle("/", app)

	srv := &http.Server{
		Handler: mux,
	}

	if cb != nil {
		cb(srv)
	}

	defer func() {
		err := srv.Close()
		if err != nil {
			app.logf("srv: %v", err)
		}
	}()

	idleConnsClosed := make(chan struct{})

	go func() {
		sigint := make(chan os.Signal, 1)
		signal.Notify(sigint, os.Interrupt, syscall.SIGTERM)
		<-sigint

		if err := srv.Shutdown(context.Background()); err != nil {
			app.logf("web.go: %v", err)
		}

		close(idleConnsClosed)
	}()

	app.logf("web.go(%d) %s", os.Getpid(), listener.Addr())

	if err := srv.Serve(listener); err != nil {
		app.logf("web.go: %v", err)
	}

	<-idleConnsClosed

	return errors.New("web.go: exit")
}

// Inspect method
func (app *Application) Inspect() string {
	return ""
}

func (app *Application) logf(format string, v ...interface{}) {

	if app.logger != nil {
		app.logger.Printf(format, v...)
	} else {
		log.Printf(format, v...)
	}
}

func (app *Application) getParams() *Params {
	ps := app.paramsPool.Get().(*Params)
	*ps = (*ps)[0:0] // reset slice
	return ps
}

func (app *Application) putParams(ps *Params) {
	if ps != nil {
		app.paramsPool.Put(ps)
	}
}

func (app *Application) recv(w http.ResponseWriter, r *http.Request) {
	if rcv := recover(); rcv != nil {
		if app.panic != nil {
			app.panic(w, r, rcv)
		} else {
			app.logf("web.go: %v", rcv)
		}
	}
}
<commit_msg>change cb to fns<commit_after>package web

import (
	"context"
	"crypto/tls"
	"errors"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

var (
	_app  *Application
	_once sync.Once
)

// Callback function
type Callback func(ctx *Context)

// PanicCallback function
type PanicCallback func(http.ResponseWriter, *http.Request, interface{})

// Application is type of a web.Application
type Application struct {
	trees       map[string]*node
	middlewares Middlewares
	logger      *log.Logger
	panic       PanicCallback
	paramsPool  sync.Pool
	maxParams   uint16
}

// Create return a singleton web.Application
func Create() *Application {
	_once.Do(func() {
		_app = newApplication()
	})
	return _app
}

// newApplication return a web.Application
func newApplication() *Application {
	app := &Application{
		middlewares: Middlewares{},
		logger:      log.New(os.Stdout, "", log.Ldate|log.Ltime),
	}
	return app
}

// SetLogger set Logger
func (app *Application) SetLogger(logger *log.Logger) {
	app.logger = logger
}

// SetPanic set Logger
func (app *Application) SetPanic(panic PanicCallback) {
	app.panic = panic
}

// Use Add the given callback function to this application.middlewares.
func (app *Application) Use(path string, callback Callback) {

	m := Middleware{
		Path:     cleanPath(path),
		Callback: callback,
	}

	app.middlewares = append(app.middlewares, m)
}

// Resource map controller path
func (app *Application) Resource(path string, controller Controller) {
	app.ResourceFn(path, controller, nil)
}

// ResourceFn map controller path and wrap fn
func (app *Application) ResourceFn(path string, controller Controller, fn func(cb Callback, keys ...string) Callback, keys ...string) {

	path = cleanPath(path)

	if fn == nil {
		app.Get(path, controller.Index)
		app.Post(path, controller.Create)
		app.Get(path+":id", controller.Detail)
		app.Patch(path+":id", controller.Update)
		app.Put(path+":id", controller.Update)
		app.Delete(path+":id", controller.Destroy)
	} else {
		app.Get(path, fn(controller.Index, keys...))
		app.Post(path, fn(controller.Create, keys...))
		app.Get(path+":id", fn(controller.Detail, keys...))
		app.Patch(path+":id", fn(controller.Update, keys...))
		app.Put(path+":id", fn(controller.Update, keys...))
		app.Delete(path+":id", fn(controller.Destroy, keys...))
	}
}

// On add event
func (app *Application) On(name string, cb Callback) {

}

// Get method
func (app *Application) Get(path string, cb Callback) {
	app.addRoute(http.MethodGet, path, cb)
}

// Post method
func (app *Application) Post(path string, cb Callback) {
	app.addRoute(http.MethodPost, path, cb)
}

// Put method
func (app *Application) Put(path string, cb Callback) {
	app.addRoute(http.MethodPut, path, cb)
}

// Patch method
func (app *Application) Patch(path string, cb Callback) {
	app.addRoute(http.MethodPatch, path, cb)
}

// Delete method
func (app *Application) Delete(path string, cb Callback) {
	app.addRoute(http.MethodDelete, path, cb)
}

// Options method
func (app *Application) Options(path string, cb Callback) {
	app.addRoute(http.MethodOptions, path, cb)
}

func (app *Application) addRoute(method, path string, cb Callback) {

	if method == "" {
		panic("method must not be empty")
	}

	if len(path) < 1 || path[0] != '/' {
		panic("path must begin with '/' in path '" + path + "'")
	}

	if cb == nil {
		panic("callback must not be nil")
	}

	if app.trees == nil {
		app.trees = make(map[string]*node)
	}

	root := app.trees[method]

	if root == nil {
		root = new(node)
		app.trees[method] = root
	}

	root.addRoute(path, cb)

	if pc := countParams(path); pc > app.maxParams {
		app.maxParams = pc
	}

	if app.paramsPool.New == nil && app.maxParams > 0 {
		app.paramsPool.New = func() interface{} {
			ps := make(Params, 0, app.maxParams)
			return &ps
		}
	}
}

// ServeHTTP
func (app *Application) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	startTime := time.Now()
	defer app.recv(w, r)
	path := r.URL.Path

	if root := app.trees[r.Method]; root != nil {

		if callback, params, tsr := root.getValue(path, app.getParams); callback != nil {

			ctx := newContext(w, r, params)
			app.putParams(params)

			app.middlewares.exec(path, ctx)

			runTime := time.Now()
			callback(ctx)
			endTime := time.Now()

			app.logf("%s %s %s %s %s", r.Method, path, endTime.Sub(startTime), runTime.Sub(startTime), endTime.Sub(runTime))

			return

		} else if r.Method != http.MethodConnect && path != "/" {

			code := http.StatusMovedPermanently

			if r.Method != http.MethodGet {
				code = http.StatusPermanentRedirect
			}

			if tsr {

				if len(path) > 1 && path[len(path)-1] == '/' {
					r.URL.Path = path[:len(path)-1]
				} else {
					r.URL.Path = path + "/"
				}

				http.Redirect(w, r, r.URL.String(), code)

				return
			}
		}
	}

	http.NotFound(w, r)
}

// ListenAndServe Serve with options on addr
func (app *Application) ListenAndServe(addr string, fns ...func(*http.Server)) error {

	l, err := net.Listen("tcp", addr)

	if err != nil {
		log.Fatal("Listen:", err)
	}

	defer l.Close()

	return app.serve(addr, l, fns...)
}

// ListenAndServeTLS Serve with tls and options on addr
func (app *Application) ListenAndServeTLS(addr string, tlsConfig *tls.Config, fns ...func(*http.Server)) error {

	l, err := tls.Listen("tcp", addr, tlsConfig)

	if err != nil {
		log.Fatal("Listen:", err)
	}

	defer l.Close()

	return app.serve(addr, l, fns...)
}

func (app *Application) serve(addr string, listener net.Listener, fns ...func(*http.Server)) error {

	mux := http.NewServeMux()

	mux.Handle("/", app)

	srv := &http.Server{
		Handler: mux,
	}

	for _, fn := range fns {
		fn(srv)
	}

	defer func() {
		err := srv.Close()
		if err != nil {
			app.logf("srv: %v", err)
		}
	}()

	idleConnsClosed := make(chan struct{})

	go func() {
		sigint := make(chan os.Signal, 1)
		signal.Notify(sigint, os.Interrupt, syscall.SIGTERM)
		<-sigint

		if err := srv.Shutdown(context.Background()); err != nil {
			app.logf("web.go: %v", err)
		}

		close(idleConnsClosed)
	}()

	app.logf("web.go(%d) %s", os.Getpid(), listener.Addr())

	if err := srv.Serve(listener); err != nil {
		app.logf("web.go: %v", err)
	}

	<-idleConnsClosed

	return errors.New("web.go: exit")
}

// Inspect method
func (app *Application) Inspect() string {
	return ""
}

func (app *Application) logf(format string, v ...interface{}) {

	if app.logger != nil {
		app.logger.Printf(format, v...)
	} else {
		log.Printf(format, v...)
	}
}

func (app *Application) getParams() *Params {
	ps := app.paramsPool.Get().(*Params)
	*ps = (*ps)[0:0] // reset slice
	return ps
}

func (app *Application) putParams(ps *Params) {
	if ps != nil {
		app.paramsPool.Put(ps)
	}
}

func (app *Application) recv(w http.ResponseWriter, r *http.Request) {
	if rcv := recover(); rcv != nil {
		if app.panic != nil {
			app.panic(w, r, rcv)
		} else {
			app.logf("web.go: %v", rcv)
		}
	}
}
<|endoftext|>