Quelle  fork.c   Sprache: C

// SPDX-License-Identifier: GPL-2.0+

/*
 * Context switch microbenchmark.
 *
 * Copyright 2018, Anton Blanchard, IBM Corp.
 */

#define _GNU_SOURCE
#include <assert.h>
#include <errno.h>
#include <getopt.h>
#include <limits.h>
#include <linux/futex.h>
#include <pthread.h>
#include <sched.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/shm.h>
#include <sys/syscall.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

static unsigned int timeout = 30;

static void set_cpu(int cpu)
{
 cpu_set_t cpuset;

 if (cpu == -1)
  return;

 CPU_ZERO(&cpuset);
 CPU_SET(cpu, &cpuset);

 if (sched_setaffinity(0, sizeof(cpuset), &cpuset)) {
  perror("sched_setaffinity");
  exit(1);
 }
}

static void start_process_on(void *(*fn)(void *), void *arg, int cpu)
{
 int pid;

 pid = fork();
 if (pid == -1) {
  perror("fork");
  exit(1);
 }

 if (pid)
  return;

 set_cpu(cpu);

 fn(arg);

 exit(0);
}

static int cpu;
static int do_fork = 0;
static int do_vfork = 0;
static int do_exec = 0;
static char *exec_file;
static int exec_target = 0;
static unsigned long iterations;
static unsigned long iterations_prev;

static void run_exec(void)
{
 char *const argv[] = { "./exec_target", NULL };

 if (execve("./exec_target", argv, NULL) == -1) {
  perror("execve");
  exit(1);
 }
}

static void bench_fork(void)
{
 while (1) {
  pid_t pid = fork();
  if (pid == -1) {
   perror("fork");
   exit(1);
  }
  if (pid == 0) {
   if (do_exec)
    run_exec();
   _exit(0);
  }
  pid = waitpid(pid, NULL, 0);
  if (pid == -1) {
   perror("waitpid");
   exit(1);
  }
  iterations++;
 }
}

static void bench_vfork(void)
{
 while (1) {
  pid_t pid = vfork();
  if (pid == -1) {
   perror("fork");
   exit(1);
  }
  if (pid == 0) {
   if (do_exec)
    run_exec();
   _exit(0);
  }
  pid = waitpid(pid, NULL, 0);
  if (pid == -1) {
   perror("waitpid");
   exit(1);
  }
  iterations++;
 }
}

static void *null_fn(void *arg)
{
 pthread_exit(NULL);
}

static void bench_thread(void)
{
 pthread_t tid;
 cpu_set_t cpuset;
 pthread_attr_t attr;
 int rc;

 rc = pthread_attr_init(&attr);
 if (rc) {
  errno = rc;
  perror("pthread_attr_init");
  exit(1);
 }

 if (cpu != -1) {
  CPU_ZERO(&cpuset);
  CPU_SET(cpu, &cpuset);

  rc = pthread_attr_setaffinity_np(&attr, sizeof(cpu_set_t), &cpuset);
  if (rc) {
   errno = rc;
   perror("pthread_attr_setaffinity_np");
   exit(1);
  }
 }

 while (1) {
  rc = pthread_create(&tid, &attr, null_fn, NULL);
  if (rc) {
   errno = rc;
   perror("pthread_create");
   exit(1);
  }
  rc = pthread_join(tid, NULL);
  if (rc) {
   errno = rc;
   perror("pthread_join");
   exit(1);
  }
  iterations++;
 }
}

static void sigalrm_handler(int junk)
{
 unsigned long i = iterations;

 printf("%ld\n", i - iterations_prev);
 iterations_prev = i;

 if (--timeout == 0)
  kill(0, SIGUSR1);

 alarm(1);
}

static void sigusr1_handler(int junk)
{
 exit(0);
}

static void *bench_proc(void *arg)
{
 signal(SIGALRM, sigalrm_handler);
 alarm(1);

 if (do_fork)
  bench_fork();
 else if (do_vfork)
  bench_vfork();
 else
  bench_thread();

 return NULL;
}

static struct option options[] = {
 { "fork", no_argument, &do_fork, 1 },
 { "vfork", no_argument, &do_vfork, 1 },
 { "exec", no_argument, &do_exec, 1 },
 { "timeout", required_argument, 0, 's' },
 { "exec-target", no_argument, &exec_target, 1 },
 { NULL },
};

static void usage(void)
{
 fprintf(stderr, "Usage: fork CPU\n\n");
 fprintf(stderr, "\t\t--fork\tUse fork() (default threads)\n");
 fprintf(stderr, "\t\t--vfork\tUse vfork() (default threads)\n");
 fprintf(stderr, "\t\t--exec\tAlso exec() (default no exec)\n");
 fprintf(stderr, "\t\t--timeout=X\tDuration in seconds to run (default 30)\n");
 fprintf(stderr, "\t\t--exec-target\tInternal option for exec workload\n");
}

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

 while (1) {
  int option_index = 0;

  c = getopt_long(argc, argv, "", options, &option_index);

  if (c == -1)
   break;

  switch (c) {
  case 0:
   if (options[option_index].flag != 0)
    break;

   usage();
   exit(1);
   break;

  case 's':
   timeout = atoi(optarg);
   break;

  default:
   usage();
   exit(1);
  }
 }

 if (do_fork && do_vfork) {
  usage();
  exit(1);
 }
 if (do_exec && !do_fork && !do_vfork) {
  usage();
  exit(1);
 }

 if (do_exec) {
  char *dirname = strdup(argv[0]);
  int i;
  i = strlen(dirname) - 1;
  while (i) {
   if (dirname[i] == '/') {
    dirname[i] = '\0';
    if (chdir(dirname) == -1) {
     perror("chdir");
     exit(1);
    }
    break;
   }
   i--;
  }
 }

 if (exec_target) {
  exit(0);
 }

 if (((argc - optind) != 1)) {
  cpu = -1;
 } else {
  cpu = atoi(argv[optind++]);
 }

 if (do_exec)
  exec_file = argv[0];

 set_cpu(cpu);

 printf("Using ");
 if (do_fork)
  printf("fork");
 else if (do_vfork)
  printf("vfork");
 else
  printf("clone");

 if (do_exec)
  printf(" + exec");

 printf(" on cpu %d\n", cpu);

 /* Create a new process group so we can signal everyone for exit */
 setpgid(getpid(), getpid());

 signal(SIGUSR1, sigusr1_handler);

 start_process_on(bench_proc, NULL, cpu);

 while (1)
  sleep(3600);

 return 0;
}