// The function timer which displays results all three minutes
func (hosts *Hosts) WorkingTimer(step int) {
// set a timer
timer := time.NewTicker(time.Duration(step) * time.Second)
// infinite loop
for {
// wait for timer
<-timer.C
// get the current time
formatter.ColoredPrintln(
formatter.Magenta,
false,
"#################################################",
)
formatter.ColoredPrintln(
formatter.Blue,
false,
"At time ",
time.Now().Format("15:04:05"),
)
// display working hosts
hosts.DisplayWorkingHosts()
}
}
// Function which executes commands when a host has to wait for other hosts.
func (host *Host) Waiter() {
// display
formatter.ColoredPrintln(
formatter.Magenta,
true,
"Waiting more jobs for", host.Hostname,
)
// say it is not working
host.IsWorking = false
// Now wait for new job
<-*(host.Wait)
// display
formatter.ColoredPrintln(
formatter.Magenta,
true,
host.Hostname, "has more jobs !",
)
formatter.ColoredPrintln(
formatter.Green,
true,
"Number of commands for", host.Hostname, ":",
len(host.Commands),
)
// indicate that it is working
host.IsWorking = true
}
// This function loop over hosts to launch commands in concurrent mode.
func (d *Dispatcher) RunCommands(
ncommands int,
) {
// number of hosts
nhosts := len(d.Hosts.Hosts)
// check that there is some hosts
if nhosts == 0 {
formatter.ColoredPrintln(
formatter.Red,
false,
"There is no hosts given to run commands !",
)
}
// A channel to wait for end of program
d.ender = make(chan bool, ncommands)
// run the routine to manage disconnections
go d.Disconnection()
// if we use the timer, run the go routine
if d.Config.Timer {
go d.Hosts.RemainingCommands(
d.Config,
ncommands,
)
}
// if we use the timer of working hosts, run the go routine
if d.Config.WorkTimer {
go d.Hosts.WorkingTimer(
d.Config.WorkTime,
)
}
// loop over hosts and run the command
for _, host := range d.Hosts.Hosts {
// in several goroutine
go d.RunOnHost(host)
}
// Wait for the end of goroutines
for i := 0; i < ncommands; i++ {
<-d.ender
formatter.ColoredPrintln(
formatter.White,
true,
"Number of remaining commands:", ncommands-1-i,
)
}
// display the summary of commands on hosts
d.Hosts.DisplayHostsCommands()
}
// This function sets a timer and check the remaining number of commands
// to execute.
func (hosts *Hosts) RemainingCommands(
config config,
ncommands int,
) {
// display
formatter.ColoredPrintln(
formatter.White,
true,
"Timer for commands is set !",
)
// set a timer
timer := time.NewTicker(time.Duration(config.GetTimer()) * time.Second)
// start an infinite loop
for {
// Wait for the timer
<-timer.C
// counter
counter := 0
// loop over hosts
for _, host := range hosts.Hosts {
// check that it is connected
if host.Connected {
// increment counter
counter += (host.CommandNumber - 1)
}
}
// display the result
formatter.ColoredPrintln(
formatter.Magenta,
false,
"Number of commands executed :",
counter,
"/",
ncommands,
)
// get the current time
hour, minute, second := time.Now().Clock()
formatter.ColoredPrintln(
formatter.Magenta,
false,
"at time ",
strconv.Itoa(hour)+":"+
strconv.Itoa(minute)+":"+strconv.Itoa(second),
)
}
}
// load a private key
func loadPEM(file string) ([]byte, error) {
// open the file
f, err := os.Open(file)
defer func() {
err := f.Close()
if err != nil {
formatter.ColoredPrintln(
formatter.Red,
false,
"The file can't be closed for the private key!\n",
"Reason is: ", err.Error(),
)
}
}()
// check errors when opening
if err != nil {
return nil, err
}
// read data
buf := bytes.NewBuffer(nil)
_, err = io.Copy(buf, f)
if err != nil {
return nil, err
}
// parse private keys
return buf.Bytes(), nil
}
// expanduser
func Expanduser(path string) string {
var home string
// get the current user
if usr, err := user.Current(); err == nil {
// get the home directory of the current user
home = usr.HomeDir
} else {
// an error occurred, fallback to the home variable
home = os.ExpandEnv("$HOME")
}
// check the path in input
if len(path) < 1 {
formatter.ColoredPrintln(
formatter.Red,
false,
"The length of the path isn't sufficient!",
)
}
// replace the tilde by home
if path[:1] == "~" {
path = strings.Replace(path, "~", home, 1)
}
return path
}
// A function to read an hosts file in the YAML format and returns
// a dictionary in the same format as the structured file.
func ReadHostsYAML(
filename string,
) *Hosts {
// Start by reading the whole file in byte
data, _ := ioutil.ReadFile(filename)
// Create the variable handling the type of the user file
t := &Hosts{}
// Now read in the YAML file the structure of the file into
// the structured dictionary
err := goyaml.Unmarshal(
data,
t,
)
// Check error when reading the file
if err != nil {
formatter.ColoredPrintln(
formatter.Red,
false,
"The file "+filename+" can't be read for accessing"+
"the YAML structure!\n"+
"Reason is: "+err.Error(),
)
return nil
}
// return the structured file and data
return t
}
// Functions to write the results of command into a log file in a given
// directory specified in argument of the program with the option -log_command.
func WriteLogCommand(
output string,
config config,
hostname string,
command string,
number int,
) {
// Create the file name with hostname and the number of the command
filename := config.GetLogCommand() + "/" +
hostname + strconv.Itoa(number) + ".log"
// content of the file
content := command + "\n" + output
// open the file
f, err := os.Create(filename)
// defer the close
defer func() {
if err := f.Close(); err != nil {
formatter.ColoredPrintln(
formatter.Red,
false,
"The file "+filename+" can't be closed!\n"+
"Reason is: "+err.Error(),
)
}
}()
// error
if err != nil {
formatter.ColoredPrintln(
formatter.Red,
false,
"The file "+filename+" can't be open for logging!\n"+
"Reason is: "+err.Error(),
)
}
// write the content in the file
f.WriteString(content)
}
// A method for the construction of the configuration
// object necessary for the connection to the host.
func (s *Session) NewConfig(user user) error {
// get the content of the private key file
key, err := loadPEM(os.ExpandEnv(tools.Expanduser(user.GetPrivateKey())))
if err != nil {
return err
}
// parse the key
parsed, err := ssh.ParseRawPrivateKey(key)
if err != nil {
formatter.ColoredPrintln(
formatter.Red,
false,
"Can't parse the private key!\n",
"Reason is: ", err.Error(),
)
}
// convert into signer
signer, err := ssh.NewSignerFromKey(parsed)
if err != nil {
formatter.ColoredPrintln(
formatter.Red,
false,
"Can't create signer from private key!\n",
"Reason is: ", err.Error(),
)
}
// Construct the configuration with password authentication
s.Config = &ssh.ClientConfig{
User: user.GetUsername(),
Auth: []ssh.AuthMethod{
ssh.PublicKeys(signer),
},
}
return nil
}
// A function which displays hosts and the number of commands they have executed.
func (hosts *Hosts) DisplayHostsCommands() {
// counter
counter := 0
// loop over hosts
for _, host := range hosts.Hosts {
// display hostname
formatter.ColoredPrint(
formatter.Magenta,
false,
host.Hostname, ": ",
)
// display the number of command executed with different
// colors in case of a disconnected host
if host.Connected {
formatter.ColoredPrintln(
formatter.Green,
false,
host.CommandNumber, "/", len(host.Commands),
)
counter += host.CommandNumber
} else {
formatter.ColoredPrintln(
formatter.Red,
false,
host.CommandNumber-1, "/", len(host.Commands),
)
counter += host.CommandNumber - 1
}
}
// display the total number to check coherence
formatter.ColoredPrintln(
formatter.Magenta,
false,
"Total of commands:", counter,
)
}
// Function to dispatch an host on other. Set variables to allow a good synchronisation
// between go routines.
func (d *Dispatcher) Disconnection() {
for {
// display
formatter.ColoredPrintln(
formatter.Green,
false,
"Waiting for a disconnected host !",
)
// wait for a signal from a disconnected host
host := <-d.disconnected
// display
formatter.ColoredPrintln(
formatter.Green,
false,
"Dispatch the jobs of", host.Hostname,
"to other connected hosts !",
)
// mark the host as not connected
host.Connected = false
// and not working
host.IsWorking = false
// dispatch remaining work to other hosts
d.Hosts.Dispatcher(
host.Commands[host.CommandNumber-1:],
d.Config.HostsMax,
false,
)
// display
formatter.ColoredPrintln(
formatter.Green,
false,
"Dispatching done for", host.Hostname, "!",
)
}
}
// A routine to check the hosts which are executing commands at this time
func (hosts *Hosts) DisplayWorkingHosts() {
// loop over hosts
for _, host := range hosts.Hosts {
// if the host is connected and is marked as executing command
if host.Connected && host.IsWorking {
// Display the name of the host and the command being executed
formatter.ColoredPrintln(
formatter.Blue,
false,
host.Hostname, ":", host.Commands[host.CommandNumber-1].Command,
)
// Display the number of remaining commands
formatter.ColoredPrintln(
formatter.Blue,
false,
"Command",
strconv.Itoa(host.CommandNumber)+"/"+
strconv.Itoa(len(host.Commands)),
)
}
}
}
// Function to execute a disconnection of host with a command.
func (d *Dispatcher) Disconnect(
host *host.Host,
message string,
) {
// display
formatter.ColoredPrintln(
formatter.Red,
false,
message,
)
// dispatch remaining work to other hosts
d.disconnected <- host
}
// A function to return the remaining commands if the number of hosts
// available is zero
func displayRemainingCommands(commands []commands.Command) {
// Display message
formatter.ColoredPrint(
formatter.Magenta,
false,
"The list of not runned commands is:\n",
)
// loop over commands
for _, command := range commands {
// Display the command
formatter.ColoredPrintln(
formatter.Red,
false,
command.Command,
)
}
}
// Function to add a session to the connection to the host.
// Since multiple sessions can exist for a connection, we allow
// the possibility to append a session into a list of session.
// The function returns too the created session in order to
// have an easy access to the session newly created.
// TODO: Maybe add them into a dictionary in order to allow to
// use a name for retrieving the session as in tmux, etc... just by
// typing a name.
func (s *Session) AddSession() error {
// create the session
session, err := s.Client.NewSession()
// append the session to the list
if err != nil {
formatter.ColoredPrintln(
formatter.Red,
false,
"Failed to create the session to the host!\n"+
"Reason is: "+err.Error(),
)
return nil
} else {
s.Session = session
}
// return the result
return err
}
// This function dispatches the commands on some hosts
func (hosts *Hosts) Dispatcher(
commands []commands.Command,
nhosts_max int,
first bool,
) {
// the pointer to the host structure
var host int
// Determine the number of hosts available in theory
nhosts := hosts.CountConnectedHosts()
// check there is at least one host connected
if nhosts == 0 {
formatter.ColoredPrintln(
formatter.Red,
true,
"There is no hosts available to do the job !",
)
// Display not done commands
displayRemainingCommands(commands)
// Exit the program
os.Exit(1)
}
// store here the list of selected hosts
myhost := make([]int, 0)
// create the sorter
sorter := &hostSorter{
Hosts: hosts,
}
// Now sort by charge in commands
sort.Sort(sorter)
// loop over commands and affect them to hosts
host = -1
for _, command := range commands {
// pass to another host
host = (host + 1) % nhosts
// if the host isn't connected
for !hosts.Hosts[host].Connected {
// pass to the next host
host = (host + 1) % nhosts
}
// count the number of workers
nworkers := hosts.CountWorkers()
// if the maximal number of hosts is get, affect to only working
// hosts
if nworkers >= nhosts_max && nhosts_max > 0 {
for !hosts.Hosts[host].IsWorker() {
host = (host + 1) % nhosts
}
}
// add the host to list
myhost = append(myhost, host)
// append to the list of commands
hosts.Hosts[host].Commands = append(
hosts.Hosts[host].Commands,
command,
)
}
// message to say that the host has more jobs
if !first {
// loop over hosts which have more jobs
for _, host := range myhost {
// send a non blocking signal
formatter.ColoredPrintln(
formatter.None,
true,
"Send more jobs signal to", hosts.Hosts[host].Hostname,
"\nwith new length of", len(hosts.Hosts[host].Commands),
)
select {
case *(hosts.Hosts[host].Wait) <- 1:
default:
}
}
}
}
// This function is used to run a command on a host
// with supplied informations.
func (d *Dispatcher) RunOnHost(
host *host.Host,
) {
loop:
// Do an infinite loop for waiting when ended
for {
// check the size of commands to execute before
if len(host.Commands) != 0 {
// loop over commands on this hosts
for i := host.CommandNumber; i < len(host.Commands); i++ {
// display
formatter.ColoredPrintln(
formatter.Blue,
true,
"Executing command", i, "for", host.Hostname,
)
// number of the command
host.CommandNumber = i + 1
// check if we want to exclude too loaded hosts
if d.Config.ExcludeLoaded {
if is, err := host.IsTooLoaded(
host.Commands[i].User,
d.Config,
); is && err == nil {
d.Disconnect(
host,
"The host "+host.Hostname+" is too loaded!",
)
break loop
} else if err != nil {
d.Disconnect(
host,
"Problem occurred when checking loading for "+
host.Hostname+"!\n"+
"Reason is: "+err.Error(),
)
break loop
}
}
// say the host is working
host.IsWorking = true
// Execute the command on the specified host
output, err := host.OneCommand(&host.Commands[i])
if err != nil {
d.Disconnect(host, output)
break loop
}
// The command has been executed correctly, say it to other
d.ender <- true
// Write the log of the command
tools.WriteLogCommand(
output,
d.Config,
host.Hostname,
host.Commands[i].Command,
i,
)
// for now print the result of the command
if !d.Config.NoResults {
formatter.ColoredPrintln(
formatter.Magenta,
false,
output,
)
}
// wait here for new jobs
if i == len(host.Commands)-1 {
//Wait for other hosts
host.Waiter()
}
}
} else {
// Now wait for new job
host.Waiter()
}
}
}
// To test the run of commands.
func TestRunCommands(t *testing.T) {
// Read the user structure from the test file
usr, _ := myusr.Current()
users := configuration.ReadUsersYAML(usr.HomeDir + "/CONFIG/TOD/users/users.yaml")
// configuration
conf := &configuration.Config{}
conf.Port = 22
conf.Protocol = "tcp"
conf.Timeout = 10
conf.LogCommand = "/tmp"
conf.CPUMax = 25.0
conf.MemoryMax = 30.0
conf.ExcludeLoaded = true
conf.WorkTimer = true
conf.WorkTime = 120
conf.HostsMax = 5
conf.Stdin = true
// read command from the example configuration
cmds := configuration.UsersToDispatcher(*users)
// replicate the command in the example
commands := make([]commands.Command, 121)
for i := range commands {
commands[i] = cmds[0]
}
// Create the list of commands and hosts
hsts := new(host.Hosts)
hosts := make([]*host.Host, len(hostnames))
for i, hst := range hostnames {
// Create the host object in the slice
hosts[i] = &host.Host{
Hostname: hst,
}
}
hsts.Hosts = hosts
// display
formatter.ColoredPrintln(
formatter.Blue,
false,
"All data initialized !",
)
// Create dispatcher
dispatcher := New(conf, hsts)
// Dispatch commands on hosts
hsts.Dispatcher(
commands,
conf.HostsMax,
true,
)
// display
formatter.ColoredPrintln(
formatter.Blue,
false,
"Dispatching the commands on hosts done !",
)
// Run commands in concurrent
dispatcher.RunCommands(len(commands))
// display
formatter.ColoredPrintln(
formatter.Blue,
false,
"Commands done !",
)
}
