cmd/pbind/pbind.c

	pbind.c revision 7c478bd95313f5f23a4c958a745db2134aa03244
/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License, Version 1.0 only
 * (the "License").  You may not use this file except in compliance
 * with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
/*
 * Copyright 2004 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */

#pragma ident   "%Z%%M% %I% %E% SMI"

/*
 * pbind - bind a process to a processor (non-exclusively)
 */

#include <sys/types.h>
#include <sys/procset.h>
#include <sys/processor.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <procfs.h>
#include <fcntl.h>
#include <errno.h>
#include <dirent.h>
#include <locale.h>
#include <libproc.h>
#include <stdarg.h>

#if !defined(TEXT_DOMAIN)       /* should be defined by cc -D */
#define TEXT_DOMAIN     "SYS_TEST"  /* Use this only if it weren't */
#endif

#define ERR_OK      0       /* exit status for success */
#define ERR_FAIL    1       /* exit status for errors */
#define ERR_USAGE   2       /* exit status for usage errors */

static char *progname;
static char bflag;
static char qflag;
static char Qflag;
static char uflag;
static char Uflag;
static int  errors;

#define MAX_PROCFS_PATH 80

/*PRINTFLIKE1*/
static void
warn(char *format, ...)
{
    int err = errno;
    va_list alist;

    (void) fprintf(stderr, "%s: ", progname);
    va_start(alist, format);
    (void) vfprintf(stderr, format, alist);
    va_end(alist);
    if (strchr(format, '\n') == NULL)
        (void) fprintf(stderr, ": %s\n", strerror(err));
}

/*PRINTFLIKE1*/
static void
die(char *format, ...)
{
    int err = errno;
    va_list alist;

    (void) fprintf(stderr, "%s: ", progname);
    va_start(alist, format);
    (void) vfprintf(stderr, format, alist);
    va_end(alist);
    if (strchr(format, '\n') == NULL)
        (void) fprintf(stderr, ": %s\n", strerror(err));
    exit(ERR_FAIL);
}

/*
 * Output for query.
 */
static void
query_out(id_t pid, id_t lwpid, processorid_t cpu)
{
    char *proclwp;
    char pidstr[21];

    if (lwpid == -1) {
        (void) snprintf(pidstr, 20, "%d", (int)pid);
        proclwp = "process";
    } else {
        (void) snprintf(pidstr, 20, "%d/%d", (int)pid, (int)lwpid);
        proclwp = "lwp";
    }

    if (cpu == PBIND_NONE)
        (void) printf(gettext("%s id %s: not bound\n"),
            proclwp, pidstr);
    else
        (void) printf(gettext("%s id %s: %d\n"),
            proclwp, pidstr, cpu);
}

/*
 * Binding error.
 */
static void
bind_err(processorid_t cpu, id_t pid, id_t lwpid, int err)
{
    char *msg;

    switch (cpu) {
    case PBIND_NONE:
        msg = gettext("unbind");
        break;
    case PBIND_QUERY:
        msg = gettext("query");
        break;
    default:
        msg = gettext("bind");
        break;
    }
    if (lwpid == -1)
        warn(gettext("cannot %s pid %d: %s\n"), msg,
            (int)pid, strerror(err));
    else
        warn(gettext("cannot %s lwpid %d/%d: %s\n"), msg,
            (int)pid, (int)lwpid, strerror(err));
}

/*
 * Output for bind.
 */
static void
bind_out(id_t pid, id_t lwpid, processorid_t old, processorid_t new)
{
    char *proclwp;
    char pidstr[21];

    if (lwpid == -1) {
        (void) snprintf(pidstr, 20, "%d", (int)pid);
        proclwp = "process";
    } else {
        (void) snprintf(pidstr, 20, "%d/%d", (int)pid, (int)lwpid);
        proclwp = "lwp";
    }

    if (old == PBIND_NONE) {
        if (new == PBIND_NONE)
            (void) printf(gettext("%s id %s: was not bound, "
                "now not bound\n"), proclwp, pidstr);
        else
            (void) printf(gettext("%s id %s: was not bound, "
                "now %d\n"), proclwp, pidstr, new);
    } else {
        if (new == PBIND_NONE)
            (void) printf(gettext("%s id %s: was %d, "
                "now not bound\n"), proclwp, pidstr, old);
        else
            (void) printf(gettext("%s id %s: was %d, "
                "now %d\n"), proclwp, pidstr, old, new);
    }
}

static struct ps_prochandle *
grab_proc(id_t pid)
{
    int ret;
    struct ps_prochandle *Pr;

    if ((Pr = Pgrab(pid, 0, &ret)) == NULL) {
        warn(gettext("cannot control process %d: %s\n"),
            (int)pid, Pgrab_error(ret));
        errors = ERR_FAIL;
        return (NULL);
    }

    /*
     * Set run-on-last-close flag so the controlled process
     * runs even if we die on a signal, and create an agent LWP.
     */
    if (Psetflags(Pr, PR_RLC) != 0 || Pcreate_agent(Pr) != 0) {
        warn(gettext("cannot control process %d\n"), (int)pid);
        errors = ERR_FAIL;
        Prelease(Pr, 0);
        return (NULL);
    }
    return (Pr);
}

static void
rele_proc(struct ps_prochandle *Pr)
{
    if (Pr == NULL)
        return;
    Pdestroy_agent(Pr);
    Prelease(Pr, 0);
}

static void
bind_lwp(struct ps_prochandle *Pr, id_t pid, id_t lwpid, processorid_t cpu)
{
    processorid_t old_cpu;

    if (pr_processor_bind(Pr, P_LWPID, lwpid, cpu, &old_cpu) < 0) {
        bind_err(cpu, pid, lwpid, errno);
        errors = ERR_FAIL;
    } else {
        if (qflag)
            query_out(pid, lwpid, old_cpu);
        else
            bind_out(pid, lwpid, old_cpu, cpu);
    }
}

/*
 * Query, set, or clear bindings for the range of LWPs in the given process.
 */
static int
do_lwps(id_t pid, const char *range, processorid_t cpu)
{
    char procfile[MAX_PROCFS_PATH];
    struct ps_prochandle *Pr;
    struct prheader header;
    processorid_t binding;
    struct lwpsinfo *lwp;
    char *lpsinfo, *ptr;
    int nent, size;
    int i, fd, found;

    /*
     * Report bindings for LWPs in process 'pid'.
     */
    (void) snprintf(procfile, MAX_PROCFS_PATH,
        "/proc/%d/lpsinfo", (int)pid);
    if ((fd = open(procfile, O_RDONLY)) < 0) {
        if (errno == ENOENT)
            errno = ESRCH;
        bind_err(cpu, pid, -1, errno);
        return (ERR_FAIL);
    }
    if (pread(fd, &header, sizeof (header), 0) != sizeof (header)) {
        (void) close(fd);
        bind_err(cpu, pid, -1, errno);
        return (ERR_FAIL);
    }
    nent = header.pr_nent;
    size = header.pr_entsize * nent;
    ptr = lpsinfo = malloc(size);
    if (lpsinfo == NULL) {
        bind_err(cpu, pid, -1, errno);
        return (ERR_FAIL);
    }
    if (pread(fd, lpsinfo, size, sizeof (header)) != size) {
        bind_err(cpu, pid, -1, errno);
        free(lpsinfo);
        (void) close(fd);
        return (ERR_FAIL);
    }

    if ((bflag || uflag) && (Pr = grab_proc(pid)) == NULL) {
        free(lpsinfo);
        (void) close(fd);
        return (ERR_FAIL);
    }
    found = 0;
    for (i = 0; i < nent; i++, ptr += header.pr_entsize) {
        /*LINTED ALIGNMENT*/
        lwp = (lwpsinfo_t *)ptr;
        binding = lwp->pr_bindpro;
        if (!proc_lwp_in_set(range, lwp->pr_lwpid))
            continue;
        found++;
        if (bflag || uflag)
            bind_lwp(Pr, pid, lwp->pr_lwpid, cpu);
        else if (binding != PBIND_NONE)
            query_out(pid, lwp->pr_lwpid, binding);
    }
    if (bflag || uflag)
        rele_proc(Pr);
    free(lpsinfo);
    (void) close(fd);
    if (found == 0) {
        warn(gettext("cannot %s lwpid %d/%s: "
            "No matching LWPs found\n"),
            bflag ? "bind" : "query", pid, range);
        return (ERR_FAIL);
    }
    return (ERR_OK);
}

/*ARGSUSED*/
static int
query_all_proc(psinfo_t *psinfo, lwpsinfo_t *lwpsinfo, void *arg)
{
    id_t pid = psinfo->pr_pid;
    processorid_t binding;

    if (processor_bind(P_PID, pid, PBIND_QUERY, &binding) < 0) {
        /*
         * Ignore search errors.  The process may have exited
         * since we read the directory.
         */
        if (errno == ESRCH)
            return (0);
        bind_err(PBIND_QUERY, pid, -1, errno);
        errors = ERR_FAIL;
        return (0);
    }
    if (binding != PBIND_NONE)
        query_out(pid, -1, binding);
    return (0);
}

static int
query_all_lwp(psinfo_t *psinfo, lwpsinfo_t *lwpsinfo, void *arg)
{
    id_t pid = psinfo->pr_pid;
    id_t lwpid = lwpsinfo->pr_lwpid;
    processorid_t *cpuid = arg;
    processorid_t binding = lwpsinfo->pr_bindpro;

    if (psinfo->pr_nlwp == 1)
        lwpid = -1; /* report process bindings if only 1 lwp */
    if ((cpuid != NULL && *cpuid == binding) ||
        (cpuid == NULL && binding != PBIND_NONE))
        query_out(pid, lwpid, binding);
    return (0);
}

static int
usage(void)
{
    (void) fprintf(stderr,
        gettext("usage: \n\t%1$s -b processor_id pid[/lwpids] ...\n"
        "\t%1$s -U [processor_id] ...\n"
        "\t%1$s -Q [processor_id] ...\n"
        "\t%1$s -u pid[/lwpids] ...\n"
        "\t%1$s [-q] [pid[/lwpids] ...]\n"),
        progname);
    return (ERR_USAGE);
}

int
main(int argc, char *argv[])
{
    int c;
    int ret;
    id_t pid;
    processorid_t cpu, old_cpu;
    char *endstr;

    progname = argv[0]; /* put actual command name in messages */

    (void) setlocale(LC_ALL, "");   /* setup localization */
    (void) textdomain(TEXT_DOMAIN);

    while ((c = getopt(argc, argv, "b:qQuU")) != EOF) {
        switch (c) {

        case 'b':
            bflag = 1;
            cpu = strtol(optarg, &endstr, 10);
            if (endstr != NULL && *endstr != '\0' || cpu < 0)
                die(gettext("invalid processor ID %s\n"),
                    optarg);
            break;

        case 'q':
            qflag = 1;
            cpu = PBIND_QUERY;
            break;

        case 'Q':
            Qflag = 1;
            cpu = PBIND_QUERY;
            break;

        case 'u':
            uflag = 1;
            cpu = PBIND_NONE;
            break;

        case 'U':
            Uflag = 1;
            break;

        default:
            return (usage());
        }
    }


    /*
     * Make sure that at most one of the options b, q, Q, u, or U
     * was specified.
     */
    c = bflag + qflag + Qflag + uflag + Uflag;
    if (c < 1) {                /* nothing specified */
        qflag = 1;          /* default to query */
        cpu = PBIND_QUERY;
    } else if (c > 1) {
        warn(gettext("options -b, -q, -Q, -u and -U "
            "are mutually exclusive\n"));
        return (usage());
    }

    errors = 0;
    argc -= optind;
    argv += optind;

    /*
     * Handle query of all processes.
     */
    if (argc == 0) {
        if (bflag || uflag) {
            warn(gettext("must specify at least one pid\n"));
            return (usage());
        }
        if (Uflag) {
            if (processor_bind(P_ALL, 0, PBIND_NONE, &old_cpu) != 0)
                die(gettext("failed to unbind some LWPs"));
        }
        if (Qflag) {
            (void) proc_walk(query_all_lwp, NULL, PR_WALK_LWP);
            return (errors);
        } else {
            (void) proc_walk(query_all_proc, NULL, PR_WALK_PROC);
            return (errors);
        }
    }

    if (Qflag || Uflag) {
        /*
         * Go through listed processor IDs.
         */
        for (; argc > 0; argv++, argc--) {
            errno = 0;
            cpu = (id_t)strtol(*argv, &endstr, 10);
            if (errno != 0 || (endstr != NULL && *endstr != '\0') ||
                p_online(cpu, P_STATUS) == -1) {
                warn(gettext("invalid processor ID\n"));
                continue;
            }
            if (Qflag) {
                (void) proc_walk(query_all_lwp,
                    &cpu, PR_WALK_LWP);
                continue;
            }
            if (Uflag) {
                if (processor_bind(P_CPUID, cpu,
                    PBIND_NONE, &old_cpu) != 0) {
                    warn(gettext("failed to unbind from "
                        "processor %d"), (int)cpu);
                    errors = ERR_FAIL;
                }
                continue;
            }
        }
        return (errors);
    }

    /*
     * Go through listed process[/lwp_ranges].
     */
    for (; argc > 0; argv++, argc--) {
        errno = 0;
        pid = (id_t)strtol(*argv, &endstr, 10);
        if (errno != 0 ||
            (endstr != NULL && *endstr != '\0' && *endstr != '/')) {
            warn(gettext("invalid process ID: %s\n"), *argv);
            continue;
        }
        if (endstr != NULL && *endstr == '/') {
            /*
             * Handle lwp range case
             */
            const char *lwps = (const char *)(++endstr);
            if (*lwps == '\0' ||
                proc_lwp_range_valid(lwps) != 0) {
                warn(gettext("invalid lwp range "
                    "for pid %d\n"), (int)pid);
                errors = ERR_FAIL;
                continue;
            }
            if (!qflag)
                (void) proc_initstdio();
            ret = do_lwps(pid, lwps, qflag ? PBIND_QUERY : cpu);
            if (!qflag)
                (void) proc_finistdio();
            if (ret != ERR_OK)
                errors = ret;
        } else {
            /*
             * Handle whole process case.
             */
            if (processor_bind(P_PID, pid, cpu, &old_cpu) < 0) {
                bind_err(cpu, pid, -1, errno);
                errors = ERR_FAIL;
                continue;
            }
            if (qflag)
                query_out(pid, -1, old_cpu);
            else
                bind_out(pid, -1, old_cpu, cpu);
        }
    }
    return (errors);
}