Ticket #18828: collect-ps.c

#include <errno.h>
#include <getopt.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>

#include <kernel/OS.h>

char const watch_log[] = "/system/var/log/watch-pid.log";
char const error_log[] = "/system/var/log/watch-pid.err";
bigtime_t const waiting = 60 * 1000 * 1000;

typedef struct
{
    team_id last;
    team_id max;
} team_stat_t;

typedef struct
{
    char when[255];
    team_stat_t ts;
    bool can_spawn_process;
    bool can_spawn_thread;
} log_entry_t;

bool print_log_entry(log_entry_t * le);
bool watch_pid();
void print_current_team_id();
void collect_team_stats(team_stat_t * ts);
bool can_spawn_process();
bool can_spawn_thread();

struct
{
    bool watch;
    bool has_watch_log;
    bool has_error_log;

    char watch_log[1024];
    char error_log[1024];

    FILE * err_log_fd;
} collect_ps;

void print_collect_ps()
{
    printf(
        "Watch: %s\n"
        "has_watch_log: %s (%s)\n"
        "has_error_log: %s (%s)\n"
        , collect_ps.watch ? "ok" : "ko"
        , collect_ps.has_watch_log ? "ok" : "ko"
        , collect_ps.watch_log
        , collect_ps.has_error_log ? "ok" : "ko"
        , collect_ps.error_log
        );
}

void print_error(const char * fmt, ...)
{
    if(collect_ps.has_error_log)
    {
        if(stderr == collect_ps.err_log_fd)
        {
            FILE * f = fopen(collect_ps.error_log, "a");
            if(NULL == f)
            {
                collect_ps.has_error_log = false;
                print_error("Cannot open '%s': %s\n", collect_ps.error_log, strerror(errno));
                return;
            }

            collect_ps.err_log_fd = f;
        }
    }

    va_list args;
    va_start(args, fmt);
    vfprintf(collect_ps.err_log_fd, fmt, args);
    va_end(args);

    fflush(collect_ps.err_log_fd);
}

bool is_option_of_name(struct option options[], int option_index, char const name[], size_t name_size)
{
    return 0 == strncmp(name, options[option_index].name, name_size);
}

void copy_file_option(bool * has, char * dst, size_t expected)
{
    if('-' == optarg[0] && '\0' == optarg[1])
    {
        *has = false;
        dst[0] = '\0';

        return;
    }

    size_t s = strlcpy(dst, optarg, expected);
    if(s > expected)
        exit(EXIT_FAILURE);
    *has = true;
}

bool process_options(int argc, char **argv)
{
    collect_ps.watch = false;
    collect_ps.err_log_fd = stderr;

    size_t s = 0;

    collect_ps.has_watch_log = true;
    s = strlcpy(collect_ps.watch_log, watch_log, sizeof watch_log);
    if(s > sizeof watch_log)
        exit(EXIT_FAILURE);

    collect_ps.has_error_log = true;
    s = strlcpy(collect_ps.error_log, error_log, sizeof error_log);
    if(s > sizeof error_log)
        exit(EXIT_FAILURE);

    struct {
        int watch;
    } so;

    struct option options[] =
        { {"watch", no_argument, &so.watch, 'w'}
        , {"log", required_argument, 0, 'l'}
        , {"error", required_argument, 0, 'e'}
        , {0, 0, 0, 0}
        };
    int option_index = 0;

    opterr = 0;
    do
    {
        int c = getopt_long(argc, argv, "wl:e:", options, &option_index);

        if(-1 == c)
            break;

        switch(c)
        {
            case 0:
                if(is_option_of_name(options, option_index, "watch", sizeof "watch"))
                    collect_ps.watch = (bool) so.watch;
                else if(is_option_of_name(options, option_index, "log-file", sizeof "log-file"))
                    copy_file_option(&collect_ps.has_watch_log, collect_ps.watch_log, sizeof collect_ps.watch_log);
                else if(is_option_of_name(options, option_index, "error-file", sizeof "error-file"))
                    copy_file_option(&collect_ps.has_error_log, collect_ps.error_log, sizeof collect_ps.error_log);

                break;

            case 'l':
                copy_file_option(&collect_ps.has_watch_log, collect_ps.watch_log, sizeof collect_ps.watch_log);
                break;

            case 'e':
                copy_file_option(&collect_ps.has_error_log, collect_ps.error_log, sizeof collect_ps.error_log);
                break;

            case 'w':
                collect_ps.watch = true;
                break;

            case '?':
                fprintf(stderr, "Unsupported command '%c' '%s'\n", optopt, optarg ? optarg : "(null)");
                break;

            default:
                return false;
        }
    } while(true);

    return true;
}

int main(int argc, char * argv[])
{
    process_options(argc, argv);

    if(collect_ps.watch)
        return watch_pid()
            ? EXIT_SUCCESS
            : EXIT_FAILURE;

    print_current_team_id();
    return EXIT_SUCCESS;
}

bool can_spawn_process()
{
    pid_t pid = fork();
    if(0 == pid)
        // in child, just bail
        exit(EXIT_SUCCESS);

    if(-1 == pid)
    {
        print_error("Failing to spawn child: %s\n", strerror(errno));
        return false;
    }

    return true;
}

int32 thread_dud(void * d)
{
    (void) d;
    return 0;
}

bool can_spawn_thread()
{
    thread_id tid = spawn_thread
        ( &thread_dud
        , "probe thread"
        , 0
        , NULL
        );

    if(B_NO_MORE_THREADS == tid || B_NO_MEMORY == tid)
        return false;

    status_t ret = resume_thread(tid);
    return B_OK == ret;
}

void print_current_team_id()
{
    team_info ti;
    bzero(&ti, sizeof ti);
    get_team_info(0, &ti);

    printf("%d\n", ti.team);
}

void collect_team_stats(team_stat_t * ts)
{
    int32 cookie = 0;
    team_info ti;

    bzero(&ti, sizeof ti);
    bzero(ts, sizeof *ts);

    while(B_OK == get_next_team_info(&cookie, &ti))
    {
        ts->last = ti.team;
        ts->max = MAX(ts->last, ts->max);
    }
}

bool watch_pid()
{
    do
    {
        log_entry_t le;
        collect_team_stats(&le.ts);

        time_t t = time(NULL);
        struct tm * now = localtime(&t);
        if(NULL == now)
        {
            print_error("Time error: %s\n", strerror(errno));
            return false;
        }

        size_t s = strftime(le.when, sizeof le.when - 1, "%Y-%m-%d-%H:%M:%S", now);
        if(s > sizeof le.when - 1)
        {
            print_error("Error on making timestamp\n");
            return false;
        }

        le.can_spawn_process = can_spawn_process();
        le.can_spawn_thread = can_spawn_thread();

        if(!print_log_entry(&le))
            return false;

        status_t ret = snooze(waiting);
        if(B_INTERRUPTED == ret)
            break;
    }
    while(true);

    return true;
}

bool print_log_entry(log_entry_t * le)
{
    FILE * f = NULL;
    if(collect_ps.has_watch_log)
    {
        f = fopen(collect_ps.watch_log, "a");
        if(NULL == f)
        {
            print_error("Cannot open '%s': %s\n", collect_ps.watch_log, strerror(errno));
            return false;
        }
    }
    else
        f = stdout;

    fprintf
        ( f
        , "%s\t% 12d\t% 12d\tprocess(%s)\tthread(%s)\n"
        , le->when
        , le->ts.max
        , le->ts.last
        , le->can_spawn_process ? "ok" : "ko"
        , le->can_spawn_thread ? "ok" : "ko"
        );

    if(collect_ps.has_watch_log)
        fclose(f);

    return true;
}