<commit_before>package unit

import (
	"fmt"
	"log"
	"strings"
	"time"

	"strconv"

	"github.com/spf13/afero"

	"github.com/innovate-technologies/Dispatch/dispatchd/config"
	"github.com/innovate-technologies/Dispatch/dispatchd/unit/state"

	etcd "github.com/coreos/etcd/client"
	"github.com/coreos/go-systemd/dbus"
	"golang.org/x/net/context"
)

const unitPath = "/var/run/dispatch/"

var (
	// Config is a pointer need to be set to the main configuration
	Config *config.ConfigurationInfo
	ctx    = context.Background()
	//EtcdAPI is the etcd keys api
	EtcdAPI etcd.KeysAPI
	// DBusConnection is the connection to the system's D-Bus
	DBusConnection DBusConnectionInterface
	// FS is the file system to be used
	FS = afero.NewOsFs()
)

// UnitInterface is the interface to a Unit
type UnitInterface interface {
	Start()
	Stop()
	Restart()
	Create()
	PutOnQueue()
	SaveOnEtcd()
	Destroy()
	LoadAndWatch()
	Watch()
	SetState(state.State)
	SetDesiredState(s state.State)
}

// Unit is a struct containing all info of a specific unit
type Unit struct {
	Name         string `json:"name"`
	Machine      string `json:"machine"`
	Template     string `json:"template,omitempty"` // is set with template name if from Template
	Global       string `json:"global,omitempty"`   // is set with global name if from global
	State        state.State
	DesiredState state.State
	Ports        []int64 `json:"ports"`
	Constraints  map[string]string
	UnitContent  string `json:"unitContent"`
	onEtcd       bool
	onDisk       bool
}

// GetAll returns all units in our zone
func GetAll() ([]Unit, error) {
	setUpEtcd()
	response, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/units/%s", Config.Zone), &etcd.GetOptions{})
	if err != nil {
		return nil, err
	}

	units := []Unit{}

	for _, unitNode := range response.Node.Nodes {
		pathParts := strings.Split(unitNode.Key, "/")
		units = append(units, NewFromEtcd(pathParts[len(pathParts)-1]))
	}

	return units, nil
}

// New returns a new Unit
func New() Unit {
	setUpDBus()
	return Unit{}
}

// NewFromEtcd creates a new unit with info from etcd
func NewFromEtcd(name string) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}

	setUpEtcd()
	unit := New()
	unit.onEtcd = true
	unit.Name = getKeyFromEtcd(name, "name")
	unit.Machine = getKeyFromEtcd(name, "machine")
	unit.Template = getKeyFromEtcd(name, "template")
	unit.Global = getKeyFromEtcd(name, "global")
	unit.UnitContent = getKeyFromEtcd(name, "unit")
	unit.DesiredState = state.ForString(getKeyFromEtcd(name, "desiredState"))
	unit.State = state.ForString(getKeyFromEtcd(name, "state"))

	unit.Ports = []int64{}
	portsStringArray := strings.Split(getKeyFromEtcd(name, "ports"), ",")
	for _, portString := range portsStringArray {
		port, err := strconv.ParseInt(portString, 10, 64)
		if err == nil {
			unit.Ports = append(unit.Ports, port)
		}
	}

	return unit
}

// Start starts the specific unit
func (unit *Unit) Start() {
	log.Println("Starting unit", unit.Name)
	unit.SetState(state.Starting)
	c := make(chan string)
	_, err := DBusConnection.StartUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error starting unit", unit.Name, err)
		return
	}
	result := <-c
	if result == "done" {
		log.Println("Started unit", unit.Name)
		unit.SetState(state.Active)
	} else {
		log.Println("Failed starting unit", unit.Name)
		unit.SetState(state.Dead)
	}
}

// Stop stops the unit
func (unit *Unit) Stop() {
	log.Println("Stopping unit", unit.Name)
	c := make(chan string)
	_, err := DBusConnection.StopUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error stopping unit", unit.Name, err)
		log.Println("Killing unit", unit.Name)
		DBusConnection.KillUnit(unit.Name, 9) // the big guns!
		unit.SetState(state.Dead)
		return
	}
	result := <-c
	if result == "done" {
		log.Println("Stopped unit", unit.Name)
		unit.SetState(state.Dead)
	}
}

// Restart restarts the unit
func (unit *Unit) Restart() {
	unit.Stop()
	unit.Start()
}

// Create writes the unit to the disk
func (unit *Unit) Create() {
	if unit.Name == "" {
		log.Println("Error starting unit with no name")
		return // can't create file without a name
	}
	thisUnitPath := unitPath + unit.Name

	fileContent := []byte(unit.UnitContent)

	FS.MkdirAll(unitPath, 0755)

	FS.Remove(thisUnitPath) // make sure the old unit is gone
	file, err := FS.Create(thisUnitPath)
	if err != nil {
		panic(err)
	}
	_, err = file.Write(fileContent)
	if err != nil {
		panic(err)
	}

	c := make(chan string)
	_, stopErr := DBusConnection.StopUnit(unit.Name, "fail", c) // stop unit to make sure new one is loaded
	if stopErr == nil {
		<-c // wait on completion
	}

	unit.onDisk = true
	_, dberr := DBusConnection.LinkUnitFiles([]string{thisUnitPath}, true, true)
	fmt.Println(dberr)
	DBusConnection.Reload()
}
<commit_msg>Kill faulty units<commit_after>package unit

import (
	"fmt"
	"log"
	"strings"
	"time"

	"strconv"

	"github.com/spf13/afero"

	"github.com/innovate-technologies/Dispatch/dispatchd/config"
	"github.com/innovate-technologies/Dispatch/dispatchd/unit/state"

	etcd "github.com/coreos/etcd/client"
	"github.com/coreos/go-systemd/dbus"
	"golang.org/x/net/context"
)

const unitPath = "/var/run/dispatch/"

var (
	// Config is a pointer need to be set to the main configuration
	Config *config.ConfigurationInfo
	ctx    = context.Background()
	//EtcdAPI is the etcd keys api
	EtcdAPI etcd.KeysAPI
	// DBusConnection is the connection to the system's D-Bus
	DBusConnection DBusConnectionInterface
	// FS is the file system to be used
	FS = afero.NewOsFs()
)

// UnitInterface is the interface to a Unit
type UnitInterface interface {
	Start()
	Stop()
	Restart()
	Create()
	PutOnQueue()
	SaveOnEtcd()
	Destroy()
	LoadAndWatch()
	Watch()
	SetState(state.State)
	SetDesiredState(s state.State)
}

// Unit is a struct containing all info of a specific unit
type Unit struct {
	Name         string `json:"name"`
	Machine      string `json:"machine"`
	Template     string `json:"template,omitempty"` // is set with template name if from Template
	Global       string `json:"global,omitempty"`   // is set with global name if from global
	State        state.State
	DesiredState state.State
	Ports        []int64 `json:"ports"`
	Constraints  map[string]string
	UnitContent  string `json:"unitContent"`
	onEtcd       bool
	onDisk       bool
	etcdName     string
}

// GetAll returns all units in our zone
func GetAll() ([]Unit, error) {
	setUpEtcd()
	response, err := EtcdAPI.Get(ctx, fmt.Sprintf("/dispatch/units/%s", Config.Zone), &etcd.GetOptions{})
	if err != nil {
		return nil, err
	}

	units := []Unit{}

	for _, unitNode := range response.Node.Nodes {
		pathParts := strings.Split(unitNode.Key, "/")
		units = append(units, NewFromEtcd(pathParts[len(pathParts)-1]))
	}

	return units, nil
}

// New returns a new Unit
func New() Unit {
	setUpDBus()
	return Unit{}
}

// NewFromEtcd creates a new unit with info from etcd
func NewFromEtcd(name string) Unit {
	if !strings.HasSuffix(name, ".service") {
		name += ".service"
	}

	setUpEtcd()
	unit := New()
	unit.etcdName = name
	unit.onEtcd = true
	unit.Name = getKeyFromEtcd(name, "name")
	unit.Machine = getKeyFromEtcd(name, "machine")
	unit.Template = getKeyFromEtcd(name, "template")
	unit.Global = getKeyFromEtcd(name, "global")
	unit.UnitContent = getKeyFromEtcd(name, "unit")
	unit.DesiredState = state.ForString(getKeyFromEtcd(name, "desiredState"))
	unit.State = state.ForString(getKeyFromEtcd(name, "state"))

	unit.Ports = []int64{}
	portsStringArray := strings.Split(getKeyFromEtcd(name, "ports"), ",")
	for _, portString := range portsStringArray {
		port, err := strconv.ParseInt(portString, 10, 64)
		if err == nil {
			unit.Ports = append(unit.Ports, port)
		}
	}

	return unit
}

// Start starts the specific unit
func (unit *Unit) Start() {
	log.Println("Starting unit", unit.Name)
	unit.SetState(state.Starting)
	c := make(chan string)
	_, err := DBusConnection.StartUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error starting unit", unit.Name, err)
		return
	}
	result := <-c
	if result == "done" {
		log.Println("Started unit", unit.Name)
		unit.SetState(state.Active)
	} else {
		log.Println("Failed starting unit", unit.Name)
		unit.SetState(state.Dead)
	}
}

// Stop stops the unit
func (unit *Unit) Stop() {
	log.Println("Stopping unit", unit.Name)
	c := make(chan string)
	_, err := DBusConnection.StopUnit(unit.Name, "fail", c)
	if err != nil {
		log.Println("Error stopping unit", unit.Name, err)
		log.Println("Killing unit", unit.Name)
		DBusConnection.KillUnit(unit.Name, 9) // the big guns!
		unit.SetState(state.Dead)
		return
	}
	result := <-c
	if result == "done" {
		log.Println("Stopped unit", unit.Name)
		unit.SetState(state.Dead)
	}
}

// Restart restarts the unit
func (unit *Unit) Restart() {
	unit.Stop()
	unit.Start()
}

// Create writes the unit to the disk
func (unit *Unit) Create() {
	if unit.Name == "" {
		log.Println("Error starting unit with no name")
		if unit.etcdName != "" {
			// Faulty unit
			unit.Name = unit.etcdName
			unit.Destroy()
		}
		return // can't create file without a name
	}
	thisUnitPath := unitPath + unit.Name

	fileContent := []byte(unit.UnitContent)

	FS.MkdirAll(unitPath, 0755)

	FS.Remove(thisUnitPath) // make sure the old unit is gone
	file, err := FS.Create(thisUnitPath)
	if err != nil {
		panic(err)
	}
	_, err = file.Write(fileContent)
	if err != nil {
		panic(err)
	}

	c := make(chan string)
	_, stopErr := DBusConnection.StopUnit(unit.Name, "fail", c) // stop unit to make sure new one is loaded
	if stopErr == nil {
		<-c // wait on completion
	}

	unit.onDisk = true
	_, dberr := DBusConnection.LinkUnitFiles([]string{thisUnitPath}, true, true)
	fmt.Println(dberr)
	DBusConnection.Reload()
}

{'fraction_non_alphanumeric': 0.13818966402959004, 'fraction_numerical': 0.003287783828213295, 'mean_word_length': 3.8597104343484774, 'pattern_counts': {'":': 0, '<': 9, '<?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': '1839511', 'n_tokens_mistral': 3583, 'n_tokens_neox': 3410, 'n_words': 897}