diff options
Diffstat (limited to 'sponge.c')
-rw-r--r-- | sponge.c | 373 |
1 files changed, 373 insertions, 0 deletions
diff --git a/sponge.c b/sponge.c new file mode 100644 index 0000000..969703f --- /dev/null +++ b/sponge.c @@ -0,0 +1,373 @@ +/* + * sponge.c - read in all available info from stdin, then output it to + * file named on the command line + * + * Copyright © 2006 Tollef Fog Heen + * + * This program is free software; you can redistribute it and/or + * modify it under the terms of the GNU General Public License + * version 2 as published by the Free Software Foundation. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU + * General Public License for more details. + * + * You should have received a copy of the GNU General Public License + * along with this program; if not, write to the Free Software + * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 + * USA + * + */ + +#include <stdio.h> +#include <stdlib.h> +#include <unistd.h> +#include <sys/types.h> +#include <sys/stat.h> +/* MAX() */ +#include <sys/param.h> +#include <fcntl.h> +#include <errno.h> +#include <string.h> +#include <sys/resource.h> +/* SIZE_MAX */ +#include <stdint.h> +#include <signal.h> + +#include "physmem.c" + +#define BUFF_SIZE 8192 +#define MIN_SPONGE_SIZE BUFF_SIZE +char *tmpname = NULL; + +void usage() { + printf("sponge <file>: soak up all input from stdin and write it to <file>\n"); + exit(0); +} + +/* all the signal stuff copied from gnu sort */ + +/* The set of signals that are caught. */ +static sigset_t caught_signals; + +/* Critical section status. */ +struct cs_status { + int valid; // was bool + sigset_t sigs; +}; + +/* Enter a critical section. */ +static struct cs_status cs_enter (void) { + struct cs_status status; + status.valid = (sigprocmask(SIG_BLOCK, &caught_signals, &status.sigs) == 0); + return status; +} + +/* Leave a critical section. */ +static void cs_leave (struct cs_status status) { + if (status.valid) { + /* Ignore failure when restoring the signal mask. */ + sigprocmask(SIG_SETMASK, &status.sigs, NULL); + } +} + +static void cleanup () { + if (tmpname) { + unlink(tmpname); + } +} + +static void onexit_cleanup (void) { + struct cs_status cs = cs_enter(); + cleanup(); + cs_leave(cs); +} + +static void sighandler (int sig) { + if (! SA_NOCLDSTOP) + signal(sig, SIG_IGN); + + cleanup(); + + signal(sig, SIG_DFL); + raise(sig); +} + +/* taken from coreutils sort */ +static size_t default_sponge_size (void) { + /* Let MEM be available memory or 1/8 of total memory, whichever + is greater. */ + double avail = physmem_available(); + double total = physmem_total(); + double mem = MAX(avail, total / 8); + struct rlimit rlimit; + + /* Let SIZE be MEM, but no more than the maximum object size or + system resource limits. Avoid the MIN macro here, as it is not + quite right when only one argument is floating point. Don't + bother to check for values like RLIM_INFINITY since in practice + they are not much less than SIZE_MAX. */ + size_t size = SIZE_MAX; + if (mem < size) + size = mem; + if (getrlimit(RLIMIT_DATA, &rlimit) == 0 && rlimit.rlim_cur < size) + size = rlimit.rlim_cur; +#ifdef RLIMIT_AS + if (getrlimit(RLIMIT_AS, &rlimit) == 0 && rlimit.rlim_cur < size) + size = rlimit.rlim_cur; +#endif + + /* Leave a large safety margin for the above limits, as failure can + occur when they are exceeded. */ + size /= 2; + +#ifdef RLIMIT_RSS + /* Leave a 1/16 margin for RSS to leave room for code, stack, etc. + Exceeding RSS is not fatal, but can be quite slow. */ + if (getrlimit(RLIMIT_RSS, &rlimit) == 0 && rlimit.rlim_cur / 16 * 15 < size) + size = rlimit.rlim_cur / 16 * 15; +#endif + + /* Use no less than the minimum. */ + return MAX (size, MIN_SPONGE_SIZE); +} + +void trapsignals (void) { + ssize_t i = 0; + static int const sig[] = { + /* The usual suspects. */ + SIGALRM, SIGHUP, SIGINT, SIGPIPE, SIGQUIT, SIGTERM, +#ifdef SIGPOLL + SIGPOLL, +#endif +#ifdef SIGPROF + SIGPROF, +#endif +#ifdef SIGVTALRM + SIGVTALRM, +#endif +#ifdef SIGXCPU + SIGXCPU, +#endif +#ifdef SIGXFSZ + SIGXFSZ, +#endif + }; + int nsigs = sizeof(sig) / sizeof(sig[0]); + +#if SA_NOCLDSTOP + struct sigaction act; + + sigemptyset(&caught_signals); + for (i = 0; i < nsigs; i++) { + sigaction(sig[i], NULL, &act); + if (act.sa_handler != SIG_IGN) + sigaddset(&caught_signals, sig[i]); + } + + act.sa_handler = sighandler; + act.sa_mask = caught_signals; + act.sa_flags = 0; + + for (i = 0; i < nsigs; i++) + if (sigismember(&caught_signals, sig[i])) + sigaction(sig[i], &act, NULL); +#else + for (i = 0; i < nsigs; i++) + if (signal(sig[i], SIG_IGN) != SIG_IGN) { + signal(sig[i], sighandler); + siginterrupt (sig[i], 1); + } +#endif +} + +static void write_buff_tmp(char* buff, size_t length, FILE *fd) { + if (fwrite(buff, length, 1, fd) < 1) { + perror("error writing buffer to temporary file"); + fclose(fd); + exit(1); + } +} + +static void write_buff_tmp_finish (char* buff, size_t length, FILE *fd) { + if (length) + write_buff_tmp(buff, length, fd); + if (fflush(fd) != 0) { + perror("fflush"); + exit(1); + } +} + +static void write_buff_out (char* buff, size_t length, FILE *fd) { + if (fwrite(buff, length, 1, fd) < 1) { + perror("error writing buffer to output file"); + fclose(fd); + exit(1); + } +} + +static void copy_tmpfile (FILE *tmpfile, FILE *outfile, char *buf, size_t size) { + ssize_t i; + if (lseek(fileno(tmpfile), 0, SEEK_SET)) { + perror("could to seek to start of temporary file"); + fclose(tmpfile); + exit(1); + } + while ((i = read(fileno(tmpfile), buf, size)) > 0) { + write_buff_out(buf, i, outfile); + } + if (i == -1) { + perror("read temporary file"); + fclose(tmpfile); + exit(1); + } + if (fclose(tmpfile) != 0) { + perror("read temporary file"); + exit(1); + } + if (fclose(outfile) != 0) { + perror("error writing buffer to output file"); + exit(1); + } +} + +FILE *open_tmpfile (void) { + struct cs_status cs; + int tmpfd; + FILE *tmpfile; + mode_t mask; + char *tmpdir; + char const * const template="%s/sponge.XXXXXX"; + + trapsignals(); + cs = cs_enter(); + tmpdir = getenv("TMPDIR"); + if (tmpdir == NULL) + tmpdir = "/tmp"; + /* Subtract 2 for `%s' and add 1 for the trailing NULL. */ + tmpname=malloc(strlen(tmpdir) + strlen(template) - 2 + 1); + if (! tmpname) { + perror("failed to allocate memory"); + exit(1); + } + sprintf(tmpname, template, tmpdir); + mask=umask(077); + tmpfd = mkstemp(tmpname); + umask(mask); + atexit(onexit_cleanup); // solaris on_exit(onexit_cleanup, 0); + cs_leave(cs); + + if (tmpfd < 0) { + perror("mkstemp failed"); + exit(1); + } + tmpfile = fdopen(tmpfd, "w+"); + if (! tmpfile) { + perror("fdopen"); + exit(1); + } + return tmpfile; +} + +int main (int argc, char **argv) { + char *buf, *bufstart, *outname = NULL; + size_t bufsize = BUFF_SIZE; + size_t bufused = 0; + FILE *outfile, *tmpfile = 0; + ssize_t i = 0; + size_t mem_available = default_sponge_size(); + int tmpfile_used=0; + + if (argc > 2 || (argc == 2 && strcmp(argv[1], "-h") == 0)) { + usage(); + } + if (argc == 2) { + outname = argv[1]; + } + + tmpfile = open_tmpfile(); + bufstart = buf = malloc(bufsize); + if (!buf) { + perror("failed to allocate memory"); + exit(1); + } + while ((i = read(0, buf, bufsize - bufused)) > 0) { + bufused = bufused+i; + if (bufused == bufsize) { + if ((bufsize*2) >= mem_available) { + write_buff_tmp(bufstart, bufused, tmpfile); + bufused = 0; + tmpfile_used = 1; + } + else { + bufsize *= 2; + bufstart = realloc(bufstart, bufsize); + if (!bufstart) { + perror("failed to realloc memory"); + exit(1); + } + } + } + buf = bufstart + bufused; + } + if (i < 0) { + perror("failed to read from stdin"); + exit(1); + } + + if (outname) { + mode_t mode; + struct stat statbuf; + int exists = (lstat(outname, &statbuf) == 0); + + write_buff_tmp_finish(bufstart, bufused, tmpfile); + + /* Set temp file mode to match either + * the old file mode, or the default file + * mode for a newly created file. */ + if (exists) { + mode = statbuf.st_mode; + } + else { + mode_t mask = umask(0); + umask(mask); + mode = 0666 & ~mask; + } + if (chmod(tmpname, mode) != 0) { + perror("chmod"); + exit(1); + } + + /* If it's a regular file, or does not yet exist, + * attempt a fast rename of the temp file. */ + if (((exists && + S_ISREG(statbuf.st_mode) && + ! S_ISLNK(statbuf.st_mode) + ) || ! exists) && + rename(tmpname, outname) == 0) { + tmpname=NULL; /* don't try to cleanup tmpname */ + } + else { + /* Fall back to slow copy. */ + outfile = fopen(outname, "w"); + if (!outfile) { + perror("error opening output file"); + exit(1); + } + copy_tmpfile(tmpfile, outfile, bufstart, bufsize); + } + } + else { + if (tmpfile_used) { + write_buff_tmp_finish(bufstart, bufused, tmpfile); + copy_tmpfile(tmpfile, stdout, bufstart, bufsize); + } + else if (bufused) { + /* buffer direct to stdout, no tmpfile */ + write_buff_out(bufstart, bufused, stdout); + } + } + + return 0; +} |