github.com/wtsi-ssg/wrstat/v4@v4.5.1/cmd/cron.go (about)

     1  /*******************************************************************************
     2   * Copyright (c) 2022 Genome Research Ltd.
     3   *
     4   * Author: Sendu Bala <sb10@sanger.ac.uk>
     5   *
     6   * Permission is hereby granted, free of charge, to any person obtaining
     7   * a copy of this software and associated documentation files (the
     8   * "Software"), to deal in the Software without restriction, including
     9   * without limitation the rights to use, copy, modify, merge, publish,
    10   * distribute, sublicense, and/or sell copies of the Software, and to
    11   * permit persons to whom the Software is furnished to do so, subject to
    12   * the following conditions:
    13   *
    14   * The above copyright notice and this permission notice shall be included
    15   * in all copies or substantial portions of the Software.
    16   *
    17   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
    18   * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
    19   * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
    20   * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
    21   * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
    22   * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
    23   * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
    24   ******************************************************************************/
    25  
    26  package cmd
    27  
    28  import (
    29  	"context"
    30  	"os"
    31  	"os/exec"
    32  	"path/filepath"
    33  	"strconv"
    34  	"strings"
    35  
    36  	"github.com/adhocore/gronx"
    37  	"github.com/adhocore/gronx/pkg/tasker"
    38  	"github.com/spf13/cobra"
    39  )
    40  
    41  // options for this cmd.
    42  var crontab string
    43  var cronKill bool
    44  
    45  // cronCmd represents the cron command.
    46  var cronCmd = &cobra.Command{
    47  	Use:   "cron",
    48  	Short: "Run multi on a regular basis.",
    49  	Long: `Run multi on a regular basis.
    50  
    51  This command takes the same arguments as 'wrstat multi' and will run multi with
    52  those arguments on the given --crontab schedule.
    53  
    54  The default schedule is 8am every day.
    55  
    56  This command will just run in the foreground forever until killed. You should
    57  probably use the daemonize program to daemonize this instead.
    58  
    59  If you can run this with sudo, but don't have full root privileges yourself, you
    60  won't be able to kill the root processes yourself directly. To kill off prior
    61  invocations of cron, do 'sudo wrstsat cron --kill'.
    62  `,
    63  	Run: func(cmd *cobra.Command, args []string) {
    64  		if cronKill {
    65  			killCronProcesses()
    66  
    67  			return
    68  		}
    69  
    70  		checkMultiArgs(args)
    71  
    72  		if crontab == "" {
    73  			die("--crontab must be supplied")
    74  		}
    75  
    76  		gron := gronx.New()
    77  
    78  		if !gron.IsValid(crontab) {
    79  			die("--crontab is invalid")
    80  		}
    81  
    82  		taskr := tasker.New(tasker.Option{})
    83  		taskr.Task(crontab, func(ctx context.Context) (int, error) {
    84  			err := doMultiScheduling(args)
    85  
    86  			return 0, err
    87  		})
    88  
    89  		taskr.Run()
    90  	},
    91  }
    92  
    93  func init() {
    94  	RootCmd.AddCommand(cronCmd)
    95  
    96  	// flags specific to this sub-command
    97  	cronCmd.Flags().StringVarP(&workDir, "working_directory", "w", "", "base directory for intermediate results")
    98  	cronCmd.Flags().StringVarP(&finalDir, "final_output", "f", "", "final output directory")
    99  	cronCmd.Flags().IntVarP(&multiInodes, "inodes_per_stat", "n",
   100  		defaultInodesPerJob, "number of inodes per parallel stat job")
   101  	cronCmd.Flags().StringVar(&multiCh, "ch", "", "passed through to 'wrstat walk'")
   102  	cronCmd.Flags().StringVar(&forcedQueue, "queue", "", "force a particular queue to be used when scheduling jobs")
   103  	cronCmd.Flags().StringVarP(&quota, "quota", "q", "", "csv of gid,disk,size_quota,inode_quota")
   104  	cronCmd.Flags().StringVarP(&ownersPath, "owners", "o", "", "gid,owner csv file")
   105  	cronCmd.Flags().IntVarP(&maxMem, "max_mem", "m",
   106  		basedirRAM, "maximum MBs to reserve for any job")
   107  	cronCmd.Flags().StringVarP(&crontab, "crontab", "c",
   108  		"0 17 * * *",
   109  		"crontab describing when to run, first 5 columns only")
   110  	cronCmd.Flags().BoolVar(&cronKill, "kill", false, "kill all wrstat processes on the system")
   111  }
   112  
   113  // killCronProcesses tries to kill all 'wrstat' processes on the system.
   114  func killCronProcesses() {
   115  	exePath, err := os.Executable()
   116  	if err != nil {
   117  		die("could not get own exe: %s", err)
   118  	}
   119  
   120  	exe := filepath.Base(exePath)
   121  
   122  	cmd := exec.Command("bash", "-c", `ps ax | grep "`+exe+ //nolint: gosec
   123  		`" | grep -v grep | grep -v '\--kill' | grep -o '^[ ]*[0-9]*'`)
   124  
   125  	out, err := cmd.Output()
   126  	if err != nil {
   127  		die("could not find any %s processes: %s", exe, err)
   128  	}
   129  
   130  	pids := strings.Fields(string(out))
   131  
   132  	killPIDs(pids)
   133  }
   134  
   135  // killPIDs kills the given pids.
   136  func killPIDs(pids []string) {
   137  	killed := 0
   138  
   139  	for _, pid := range pids {
   140  		pidI, err := strconv.Atoi(pid)
   141  		if err != nil {
   142  			warn("bad pid %s: %s", pid, err)
   143  
   144  			continue
   145  		}
   146  
   147  		proc, err := os.FindProcess(pidI)
   148  		if err != nil {
   149  			warn("could not find process %d", pidI)
   150  
   151  			continue
   152  		}
   153  
   154  		err = proc.Kill()
   155  		if err != nil {
   156  			warn("could not kill pid %d: %s", pidI, err)
   157  		} else {
   158  			killed++
   159  		}
   160  	}
   161  
   162  	info("killed %d processes", killed)
   163  }