<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"
)

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
	// 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 string, verifier *sign.SignatureVerifier) (*Agent, error) {
	ttldur, err := time.ParseDuration(ttl)
	if err != nil {
		return nil, err
	}

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

	return &Agent{registry, events, machine, ttldur, 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) {
	data := struct {
		Systemd *systemd.SystemdManager
		State   *AgentState
	}{
		Systemd: a.systemd,
		State:   a.state,
	}
	return json.Marshal(data)
}

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

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

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

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

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

// Initialize pushes the local machine state to the Registry
// repeatedly until it succeeds. It returns the modification
// index of the first successful response received from etcd.
func (a *Agent) Initialize() uint64 {
	log.V(1).Infof("Initializing Agent")
	a.machine.RefreshState()

	var idx uint64
	wait := time.Second
	for {
		var err error
		if idx, err = a.registry.SetMachineState(a.machine.State(), a.ttl); err == nil {
			log.V(1).Infof("Heartbeat succeeded")
			break
		}
		log.V(1).Infof("Failed heartbeat, retrying in %v", wait)
		time.Sleep(wait)
	}

	return idx
}

// 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())
	}

	a.state.Lock()
	launched := a.state.LaunchedJobs()
	a.state.Unlock()

	// Explicitly clear heartbeats of jobs that were launched
	// locally so it doesn't confuse later state calculations
	for _, j := range launched {
		a.registry.ClearJobHeartbeat(j)
	}

	for _, j := range a.registry.GetAllJobsByMachine(bootID) {
		log.Infof("Purging Job(%s)", j.Name)
		a.ForgetJob(j.Name)
		a.ReportPayloadState(j.Name, nil)

		// 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)
		}
<commit_msg>fix(purge): Simplify graceful agent shutdown<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"
)

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
	// 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 string, verifier *sign.SignatureVerifier) (*Agent, error) {
	ttldur, err := time.ParseDuration(ttl)
	if err != nil {
		return nil, err
	}

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

	return &Agent{registry, events, machine, ttldur, 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) {
	data := struct {
		Systemd *systemd.SystemdManager
		State   *AgentState
	}{
		Systemd: a.systemd,
		State:   a.state,
	}
	return json.Marshal(data)
}

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

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

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

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

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

// Initialize pushes the local machine state to the Registry
// repeatedly until it succeeds. It returns the modification
// index of the first successful response received from etcd.
func (a *Agent) Initialize() uint64 {
	log.V(1).Infof("Initializing Agent")
	a.machine.RefreshState()

	var idx uint64
	wait := time.Second
	for {
		var err error
		if idx, err = a.registry.SetMachineState(a.machine.State(), a.ttl); err == nil {
			log.V(1).Infof("Heartbeat succeeded")
			break
		}
		log.V(1).Infof("Failed heartbeat, retrying in %v", wait)
		time.Sleep(wait)
	}

	return idx
}

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

// Purge removes the Agent's state from the Registry
func (a *Agent) Purge() {
	// Continue heartbeating the agent's machine state while attempting to
	// stop all the locally-running jobs
	purged := make(chan bool)
	go a.Heartbeat(a.ttl, purged)

	bootID := a.machine.State().BootID

	//TODO(bcwaldon): The agent should not have to ask the Registry
	// which jobs it is running in its local systemd
	for _, j := range a.registry.GetAllJobsByMachine(bootID) {
		log.Infof("Unloading Job(%s)", j.Name)
		a.UnloadJob(j.Name)
	}

	// Jobs have been stopped, the heartbeat can stop
	close(purged)

	log.V(1).Info("Removing Agent from Registry")
	if err := a.registry.RemoveMachineState(bootID); err != nil {
		log.Errorf("Failed to remove Machine %s from Registry: %s", bootID, err.Error())
	}
}

// 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)
		}

{'fraction_non_alphanumeric': 0.1213982050070855, 'fraction_numerical': 0.003897024090694379, 'mean_word_length': 3.9152640742890306, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '54484', 'n_tokens_mistral': 2997, 'n_tokens_neox': 2814, 'n_words': 884}