reniced/main.go

package main

import (
	"log"
	"os"
	"os/signal"
	"os/user"
	"strings"
	"syscall"
	"time"

	"github.com/alitto/pond"
	"github.com/sevlyar/go-daemon"
	proc "github.com/shirou/gopsutil/v3/process"
)

// Пул воркеров.
var pool *pond.WorkerPool

// Канал, в который приходят уведомления для хэндлера сигналов от траппера сигналов.
var sigChan = make(chan os.Signal, 1)

// main основная функция программы.
func main() {
	// Самое время поставить траппер сигналов.
	signal.Notify(sigChan,
		syscall.SIGINT,
		syscall.SIGTERM,
		syscall.SIGQUIT)

	go sigHandler()

	var (
		cnf Config
		err error
	)

	u, err := user.Current()

	if err != nil {
		log.Fatalf("Unable to get user info for current user: %s", err)
	}

	// Для рута и для других пользователей конфиги находятся в разных местах.
	if u.Uid != "0" {
		if os.Getenv("HOME") == "" {
			log.Fatalln("Unable to get HOME environment variable value")
		}

		cnf, err = readConf(os.Getenv("HOME") + "/.reniced.yaml")
	} else {
		cnf, err = readConf("/etc/reniced.yaml")
	}

	if err != nil {
		log.Fatalf("unable to parse config: %s", err)
	}

	if cnf.Daemon {
		cntxt := &daemon.Context{
			WorkDir: "/",
			Args:    []string{"reniced"},
		}

		if cnf.Pidfile != "" {
			cntxt.PidFileName = cnf.Pidfile
		}

		if d, err := cntxt.Reborn(); err != nil {
			log.Fatal("Unable to run: ", err)

			return
		} else if d != nil {
			return
		}

		defer cntxt.Release() //nolint: errcheck
	}

	// То, ради чего всё затевалось.
	reniced(cnf)
}

// reniced основная логика.
func reniced(cnf Config) {
	pool = pond.New(cnf.MaxWorkers, cnf.NbCapacity)

	for {
		processList, err := proc.Processes()

		if err != nil {
			log.Fatalf("Unable to get process list: %s", err)
		}

		// Не самый эффективный, зато работающий способ - простой перебор массивов.
		for _, p := range processList {
			ok := pool.TrySubmit(func() {
				if processName, err := p.Name(); err == nil {
					for _, k := range cnf.Kill {
						for _, killProcessName := range k.Name {
							if processName == killProcessName {
								switch k.Sig {
								case "kill":
									_ = p.SendSignal(syscall.SIGKILL)

									if cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s), SIG%s sent",
											processName,
											strings.ToUpper(k.Sig),
										)
									}
								case "stop":
									if s, err := p.Status(); err == nil && s[0] != "stop" {
										_ = p.SendSignal(syscall.SIGSTOP)

										if cnf.Debug {
											log.Printf(
												"Matching processName with killProcessname (%s), SIG%s sent",
												processName,
												strings.ToUpper(k.Sig),
											)
										}
									} else if err == nil && s[0] == "stop" && cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s) and it is already stopped",
											processName,
										)
									}
								case "term":
									_ = p.SendSignal(syscall.SIGTERM)

									if cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s), SIG%s sent",
											processName,
											strings.ToUpper(k.Sig),
										)
									}
								case "int":
									_ = p.SendSignal(syscall.SIGINT)

									if cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s), SIG%s sent",
											processName,
											strings.ToUpper(k.Sig),
										)
									}
								case "quit":
									_ = p.SendSignal(syscall.SIGQUIT)

									if cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s), SIG%s sent",
											processName,
											strings.ToUpper(k.Sig),
										)
									}
								case "abrt":
									_ = p.SendSignal(syscall.SIGABRT)

									if cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s), SIG%s sent",
											processName,
											strings.ToUpper(k.Sig),
										)
									}
								case "hup":
									_ = p.SendSignal(syscall.SIGHUP)

									if cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s), SIG%s sent",
											processName,
											strings.ToUpper(k.Sig),
										)
									}
								case "usr1":
									_ = p.SendSignal(syscall.SIGUSR1)

									if cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s), SIG%s sent",
											processName,
											strings.ToUpper(k.Sig),
										)
									}
								case "usr2":
									_ = p.SendSignal(syscall.SIGUSR2)

									if cnf.Debug {
										log.Printf(
											"Matching processName with killProcessname (%s), SIG%s sent",
											processName,
											strings.ToUpper(k.Sig),
										)
									}
								}
							}
						}
					}

					// Для каждого процесса извлекаем его текущий nicelevel.
					if currentNiceLevel, err := p.Nice(); err == nil {
						// Вынимаем список приоритетов.
						for _, n := range cnf.Prio {
							// Для каждого приоритета, проходимся по списку имён процессов.
							for _, name := range n.Name {
								// Достаём имя текущего процесса.
								if processName, err = p.Name(); err == nil {
									// Если имя процесса совпадает с одним из списка в конфиге, то пробуем менять ему nice.
									if processName == name {
										if currentNiceLevel != n.Nice {
											_ = syscall.Setpriority(syscall.PRIO_PROCESS, int(p.Pid), int(n.Nice))

											if cnf.Debug {
												log.Printf(
													"Set niceness for %s(%d) to %d (was %d)",
													name,
													p.Pid,
													n.Nice,
													currentNiceLevel,
												)
											}
										} else if cnf.Debug {
											log.Printf("Nicelevel for %d already set to %d", p.Pid, n.Nice)
										}
									}
								}
							}
						}
					}
				}
			})

			if !ok {
				log.Printf("Unable to add task to pool")
			}

			time.Sleep(time.Millisecond * time.Duration(cnf.CmdDelay))
		}

		time.Sleep(time.Millisecond * time.Duration(cnf.LoopDelay))
	}
}

// sigHandler хэндлер сигналов. Работает на выходе приложения. Держит INT, TERM, QUIT.
func sigHandler() {
	for {
		var s = <-sigChan
		switch s {
		case syscall.SIGINT:
			log.Print("Got SIGINT, quitting")
		case syscall.SIGTERM:
			log.Print("Got SIGTERM, quitting")
		case syscall.SIGQUIT:
			log.Print("Got SIGQUIT, quitting")

		// Заходим на новую итерацию, если у нас "неинтересный" сигнал
		default:
			continue
		}

		pool.StopAndWait()

		os.Exit(0)
	}
}

/* vim: set ft=go noet ai ts=4 sw=4 sts=4: */