ckpt.c

Path MPICH-V for MPICH the MPI Implementation

#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include <signal.h>
#include <errno.h>

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

#include "ckpt.h"

extern void rs_mode_native();
extern void rs_mode_pop();

typedef void (*fn_t)(void *);

enum {
	MAX_CALLBACKS = 1000
};
static fn_t on_preckpt[MAX_CALLBACKS];
static void *on_preckpt_arg[MAX_CALLBACKS];
static unsigned num_on_preckpt;
static fn_t on_postckpt[MAX_CALLBACKS];
static void *on_postckpt_arg[MAX_CALLBACKS];
static unsigned num_on_postckpt;
static fn_t on_restart[MAX_CALLBACKS];
static void *on_restart_arg[MAX_CALLBACKS];
static unsigned num_on_restart;

void
ckpt_on_preckpt(fn_t f, void *arg)
{
	if (num_on_preckpt >= MAX_CALLBACKS) {
		fprintf(stderr, "Warning: too many pre-ckpt callbacks\n");
		return;
	}
	on_preckpt[num_on_preckpt] = f;
	on_preckpt_arg[num_on_preckpt++] = arg;
}

void
ckpt_on_postckpt(fn_t f, void *arg)
{
	if (num_on_postckpt >= MAX_CALLBACKS) {
		fprintf(stderr, "Warning: too many post-ckpt callbacks\n");
		return;
	}
	on_postckpt[num_on_postckpt] = f;
	on_postckpt_arg[num_on_postckpt++] = arg;
}

void
ckpt_on_restart(fn_t f, void *arg)
{
	if (num_on_restart >= MAX_CALLBACKS) {
		fprintf(stderr, "Warning: too many restart callbacks\n");
		return;
	}
	on_restart[num_on_restart] = f;
	on_restart_arg[num_on_restart++] = arg;
}

static int
get_ckpt_stream()
{
	if (OPTIONS.ckpt_server)
		return request_ckpt_save(OPTIONS.ckpt_server,
					 OPTIONS.ckpt_id);
	else if (OPTIONS.ckpt_filename)
		return open(OPTIONS.ckpt_filename,
			    O_CREAT|O_TRUNC|O_WRONLY, 0600);
	else
		return open("/tmp/ckpt",
			    O_CREAT|O_TRUNC|O_WRONLY, 0600);
}

static int
save_ckpt(const struct ckpt_header *head,
	  const memregion_t *regions)
{
	int fd;
	int i;
	unsigned long total = 0;

	fd = get_ckpt_stream();
	if (0 > fd) {
		fprintf(stderr, "cannot obtain a checkpoint stream\n");
		return -1;
	}

	if (0 > xwrite(fd, head, sizeof(*head))) {
		fprintf(stderr, "head write failed\n");
		goto err;
	}
	total += sizeof(struct ckpt_header);
	if (0 > xwrite(fd, regions, head->num_regions*sizeof(*regions))) {
		fprintf(stderr, "head count failed\n");
		goto err;
	}
	total += head->num_regions*sizeof(*regions);
	for (i = 0; i < head->num_regions; i++) {
#if 0
		fprintf(stderr, "[%d] %08x - %08x\n", i,
			regions[i].addr, regions[i].addr+regions[i].len);
#endif
		if (0 > xwrite(fd, (void*)regions[i].addr, regions[i].len)) {
			fprintf(stderr, "region %d failed (%s)\n", i,
				strerror(errno));
			goto err;
		}
		total += regions[i].len;
	}
	close(fd);
	return 0;
err:
	fprintf(stderr, "Checkpoint write error\n");
	return -1;
}

static int
getcmd(char *cmd, int max)
{
	int fd;
	int rv;

	fd = open("/proc/self/cmdline", O_RDONLY);
	if (0 > fd)
		return -1;
	rv = read(fd, cmd, max);
	close(fd);
	if (0 >= rv)
		return -1;
	if (rv >= max)
		cmd[max-1] = '\0';
	else
		cmd[rv] = '\0';
	return 0;
}

int
do_ckpt()
{
	struct ckpt_header head;
	memregion_t regions[MAXREGIONS];
	int i;

	for (i = 0; i < num_on_preckpt; i++)
		on_preckpt[i](on_preckpt_arg[i]);

	if (0 > getcmd(head.cmd, sizeof(head.cmd))) {
		fprintf(stderr, "cannot read my command\n");
		return -1;
	}

	fprintf(stderr, "head.cmd is %s\n", head.cmd);

	if (0 > read_self_regions(regions, &head.num_regions)) {
		fprintf(stderr, "cannot read my memory map\n");
		return -1;
	}

	if (0 == setjmp(head.jbuf)) {
		/* Checkpoint */
		if (0 > ckpt_signals()) {
			fprintf(stderr, "cannot save the signal state\n");
			return -1;
		}
		head.brk = (unsigned long) sbrk(0);
		if (0 > save_ckpt(&head, regions)) {
			fprintf(stderr, "cannot save the ckpt image\n");
			return -1;
		}
		if (!OPTIONS.ckpt_continue)
			_exit(0); /* do not call atexit functions */
		for (i = 0; i < num_on_postckpt; i++)
			on_postckpt[i](on_postckpt_arg[i]);
		return 0;
	}		

	/* Restart */
	if (0 > unmap_ifnot_orig(regions, head.num_regions)) {
		fprintf(stderr,
			"cannot purge restart code from address space\n");
		return -1;
	}
	if (0 > set_orig_mprotect(regions, head.num_regions)) {
		fprintf(stderr, "cannot restore address space protection\n");
	}
	if (0 > restore_signals()) {
		fprintf(stderr, "cannot restore signal disposition\n");
		return -1;
	}
	for (i = num_on_restart-1; i >= 0; i--)
		on_restart[i](on_restart_arg[i]);
	return 0;
}

static void
ckpt_sig_handler(int signum)
{
	do_ckpt();
}

static void
ckpt_banner()
{
	return;
	fprintf(stderr, "libckpt of %s %s\n", __DATE__, __TIME__);
}

static void
install_ckpt_handler()
{
	if (SIG_ERR == signal(OPTIONS.ckpt_signal, ckpt_sig_handler))
		fprintf(stderr, "cannot install checkpoint signal handler\n");
}

void
ckpt_init()
{
	ckpt_banner();
	get_options();
	if (OPTIONS.ckpt_server) {
		/* FIXME: This is necessary only when interoperability
		   is disabled in the ckpting process; otherwise,
		   we'll fall back upon connection. */
		call_if_present("rs_mode_native", "librocks.so");
		if (0 <= request_ckpt_access(OPTIONS.ckpt_server,
					     OPTIONS.ckpt_id))
			/* we've been ckpted before */
			restart();
		call_if_present("rs_mode_pop", "librocks.so");
	}
	install_ckpt_handler();
}