/*      $NetBSD: schedctl.c,v 1.16 2014/07/27 04:46:48 dholland Exp $   */

/*
* Copyright (c) 2008, Mindaugas Rasiukevicius <rmind at NetBSD org>
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions
* are met:
* 1. Redistributions of source code must retain the above copyright
*    notice, this list of conditions and the following disclaimer.
* 2. Redistributions in binary form must reproduce the above copyright
*    notice, this list of conditions and the following disclaimer in the
*    documentation and/or other materials provided with the distribution.
*
* THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
* ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
* OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
* HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
* LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
* OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
* SUCH DAMAGE.
*/

/*
* schedctl(8) - a program to control scheduling of processes and threads.
*/

#include <sys/cdefs.h>

#ifndef lint
__RCSID("$NetBSD: schedctl.c,v 1.16 2014/07/27 04:46:48 dholland Exp $");
#endif

#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <err.h>
#include <fcntl.h>
#include <kvm.h>
#include <unistd.h>

#include <sys/pset.h>
#include <sys/sched.h>
#include <sys/sysctl.h>
#include <sys/types.h>

static const char *class_str[] = {
       "SCHED_OTHER",
       "SCHED_FIFO",
       "SCHED_RR",
       NULL
};

static void     sched_set(pid_t, lwpid_t, int, struct sched_param *, cpuset_t *);
static void     thread_info(pid_t, lwpid_t);
static cpuset_t *makecpuset(char *);
static void     printcpuset(cpuset_t *);
__dead static void      usage(void);

static u_int    ncpu;

int
main(int argc, char **argv)
{
       kvm_t *kd;
       struct kinfo_lwp *lwp_list, *lwp;
       struct sched_param *sp;
       cpuset_t *cpuset;
       int i, count, ch, policy;
       pid_t pid;
       lwpid_t lid;
       bool set;

       ncpu = sysconf(_SC_NPROCESSORS_CONF);

       pid = lid = 0;
       cpuset = NULL;
       set = false;

       sp = calloc(1, sizeof(struct sched_param));
       if (sp == NULL)
               err(EXIT_FAILURE, "calloc");

       sp->sched_priority = PRI_NONE;
       policy = SCHED_NONE;

       while ((ch = getopt(argc, argv, "A:C:P:p:t:")) != -1) {
               switch (ch) {
               case 'p':
                       /* PID */
                       pid = atoi(optarg);
                       break;
               case 't':
                       /* Thread (LWP) ID */
                       lid = atoi(optarg);
                       break;
               case 'A':
                       /* Affinity */
                       cpuset = makecpuset(optarg);
                       if (cpuset == NULL) {
                               fprintf(stderr, "%s: invalid CPU value\n",
                                   getprogname());
                               exit(EXIT_FAILURE);
                       }
                       break;
               case 'C':
                       /* Scheduling class */
                       for (policy = 0; class_str[policy] != NULL; policy++) {
                               if (strcasecmp(optarg, class_str[policy]) == 0)
                                       break;
                       }
                       if (class_str[policy] == NULL)
                               policy = atoi(optarg);
                       if (policy < SCHED_OTHER || policy > SCHED_RR) {
                               fprintf(stderr,
                                   "%s: invalid scheduling class\n",
                                   getprogname());
                               exit(EXIT_FAILURE);
                       }
                       set = true;
                       break;
               case 'P':
                       /* Priority */
                       sp->sched_priority = atoi(optarg);
                       if (sp->sched_priority < sysconf(_SC_SCHED_PRI_MIN) ||
                           sp->sched_priority > sysconf(_SC_SCHED_PRI_MAX)) {
                               fprintf(stderr, "%s: invalid priority\n",
                                   getprogname());
                               exit(EXIT_FAILURE);
                       }
                       set = true;
                       break;
               default:
                       usage();
               }
       }

       /* At least PID must be specified */
       if (pid == 0) {
               if (argv[optind] == NULL || lid != 0)
                       usage();
               pid = getpid();
       } else {
               if (argv[optind] != NULL)
                       usage();
       }

       /* Set the scheduling information for thread/process */
       sched_set(pid, lid, policy, set ? sp : NULL, cpuset);

       /* Show information about each thread */
       if (pid != getpid()) {
               kd = kvm_open(NULL, NULL, NULL, KVM_NO_FILES, "kvm_open");
               if (kd == NULL)
                       err(EXIT_FAILURE, "kvm_open");
               lwp_list = kvm_getlwps(kd, pid, 0, sizeof(struct kinfo_lwp), &count);
               if (lwp_list == NULL)
                       err(EXIT_FAILURE, "kvm_getlwps");
               for (lwp = lwp_list, i = 0; i < count; lwp++, i++) {
                       if (lid && lid != lwp->l_lid)
                               continue;
                       if (lwp->l_stat == LSIDL || lwp->l_stat == LSZOMB)
                               continue;
                       thread_info(pid, lwp->l_lid);
               }
               kvm_close(kd);
               free(sp);
               cpuset_destroy(cpuset);
               return 0;
       }

       (void)execvp(argv[optind], argv + optind);
       err(EXIT_FAILURE, "execvp");
}

static void
sched_set(pid_t pid, lwpid_t lid, int policy,
   struct sched_param *sp, cpuset_t *cpuset)
{
       int error;

       if (sp) {
               /* Set the scheduling parameters for the thread */
               error = _sched_setparam(pid, lid, policy, sp);
               if (error < 0)
                       err(EXIT_FAILURE, "_sched_setparam");
       }
       if (cpuset) {
               /* Set the CPU-set for affinity */
               error = _sched_setaffinity(pid, lid,
                   cpuset_size(cpuset), cpuset);
               if (error < 0)
                       err(EXIT_FAILURE, "_sched_setaffinity");
       }
}

static void
thread_info(pid_t pid, lwpid_t lid)
{
       struct sched_param sp;
       cpuset_t *cpuset;
       int error, policy;

       cpuset = cpuset_create();
       if (cpuset == NULL)
               err(EXIT_FAILURE, "cpuset_create");

       error = _sched_getparam(pid, lid, &policy, &sp);
       if (error < 0)
               err(EXIT_FAILURE, "_sched_getparam");

       error = _sched_getaffinity(pid, lid, cpuset_size(cpuset), cpuset);
       if (error < 0)
               err(EXIT_FAILURE, "_sched_getaffinity");

       printf("  LID:              %d\n", lid);
       printf("  Priority:         %d\n", sp.sched_priority);
       printf("  Class:            %s\n", class_str[policy]);

       printf("  Affinity (CPUs):  ");
       printcpuset(cpuset);
       printf("\n");

       cpuset_destroy(cpuset);
}

static cpuset_t *
makecpuset(char *str)
{
       cpuset_t *cpuset;
       char *cpustr, *s;

       if (str == NULL)
               return NULL;

       cpuset = cpuset_create();
       if (cpuset == NULL)
               err(EXIT_FAILURE, "cpuset_create");
       cpuset_zero(cpuset);

       cpustr = strdup(str);
       if (cpustr == NULL)
               err(EXIT_FAILURE, "strdup");
       s = cpustr;

       while (s != NULL) {
               char *p;
               int i;

               /* Get the CPU number and validate the range */
               p = strsep(&s, ",");
               if (p == NULL) {
                       cpuset_destroy(cpuset);
                       cpuset = NULL;
                       break;
               }
               i = atoi(p);
               if (i == -1) {
                       cpuset_zero(cpuset);
                       break;
               }
               if ((unsigned int)i >= ncpu) {
                       cpuset_destroy(cpuset);
                       cpuset = NULL;
                       break;
               }

               /* Set the bit */
               cpuset_set(i, cpuset);
       }

       free(cpustr);
       return cpuset;
}

static void
printcpuset(cpuset_t *cpuset)
{
       unsigned int i;
       bool seen;

       seen = false;
       for (i = 0; i < ncpu; i++) {
               if (cpuset_isset(i, cpuset)) {
                       if (seen) {
                               putchar(',');
                       }
                       printf("%d", i);
                       seen = true;
               }
       }

       if (!seen) {
               printf("<none>");
       }
}

static void
usage(void)
{

       fprintf(stderr, "usage: %s [-A processor] [-C class] "
           "[-P priority] [-t lid] {-p pid|command}\n", getprogname());
       exit(EXIT_FAILURE);
}