sponge.c

   1 /*
   2  *  sponge.c - read in all available info from stdin, then output it to
   3  *  file named on the command line
   4  *
   5  *  Copyright ©  2006  Tollef Fog Heen
   6  *
   7  *  This program is free software; you can redistribute it and/or
   8  *  modify it under the terms of the GNU General Public License
   9  *  version 2 as published by the Free Software Foundation.
  10  *
  11  *  This program is distributed in the hope that it will be useful,
  12  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  *  General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU General Public License
  17  * along with this program; if not, write to the Free Software
  18  * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
  19  * USA
  20  *
  21  */
  22
  23 #include <stdio.h>
  24 #include <stdlib.h>
  25 #include <unistd.h>
  26 #include <sys/types.h>
  27 #include <sys/stat.h>
  28 /* MAX() */
  29 #include <sys/param.h>
  30 #include <fcntl.h>
  31 #include <errno.h>
  32 #include <string.h>
  33 #include <sys/resource.h>
  34 /* SIZE_MAX */
  35 #include <stdint.h>
  36 #include <signal.h>
  37
  38 #include "physmem.c"
  39
  40 #define BUFF_SIZE           8192
  41 #define MIN_SPONGE_SIZE     BUFF_SIZE
  42 #define DEFAULT_TMP_NAME    "/tmp/sponge.XXXXXX"
  43 char tmpname[] = DEFAULT_TMP_NAME;
  44
  45 void usage() {
  46         printf("sponge <file>: soak up all input from stdin and write it to <file>\n");
  47         exit(0);
  48 }
  49
  50 /* all the signal stuff copied from gnu sort */
  51
  52 /* The set of signals that are caught.  */
  53 static sigset_t caught_signals;
  54
  55 /* Critical section status.  */
  56 struct cs_status {
  57         int valid; // was bool
  58         sigset_t sigs;
  59 };
  60
  61 /* Enter a critical section.  */
  62 static struct cs_status cs_enter (void) {
  63         struct cs_status status;
  64         status.valid = (sigprocmask(SIG_BLOCK, &caught_signals, &status.sigs) == 0);
  65         return status;
  66 }
  67
  68 /* Leave a critical section.  */
  69 static void cs_leave (struct cs_status status) {
  70         if (status.valid) {
  71                 /* Ignore failure when restoring the signal mask. */
  72                 sigprocmask(SIG_SETMASK, &status.sigs, NULL);
  73         }
  74 }
  75
  76 static void cleanup() {
  77         if (strcmp(tmpname, DEFAULT_TMP_NAME)) {
  78                 unlink(tmpname);
  79         }
  80 }
  81
  82 static void onexit_cleanup (void) {
  83         struct cs_status cs = cs_enter();
  84         cleanup();
  85         cs_leave(cs);
  86 }
  87
  88 static void sighandler (int sig) {
  89         if (! SA_NOCLDSTOP)
  90                 signal(sig, SIG_IGN);
  91
  92         cleanup();
  93
  94         signal(sig, SIG_DFL);
  95         raise(sig);
  96 }
  97
  98 /* taken from coreutils sort */
  99 static size_t default_sponge_size (void) {
 100         /* Let MEM be available memory or 1/8 of total memory, whichever
 101            is greater.  */
 102         double avail = physmem_available();
 103         double total = physmem_total();
 104         double mem = MAX(avail, total / 8);
 105         struct rlimit rlimit;
 106
 107         /* Let SIZE be MEM, but no more than the maximum object size or
 108            system resource limits.  Avoid the MIN macro here, as it is not
 109            quite right when only one argument is floating point.  Don't
 110            bother to check for values like RLIM_INFINITY since in practice
 111            they are not much less than SIZE_MAX.  */
 112         size_t size = SIZE_MAX;
 113         if (mem < size)
 114                 size = mem;
 115         if (getrlimit(RLIMIT_DATA, &rlimit) == 0 && rlimit.rlim_cur < size)
 116                 size = rlimit.rlim_cur;
 117 #ifdef RLIMIT_AS
 118         if (getrlimit(RLIMIT_AS, &rlimit) == 0 && rlimit.rlim_cur < size)
 119                 size = rlimit.rlim_cur;
 120 #endif
 121
 122         /* Leave a large safety margin for the above limits, as failure can
 123            occur when they are exceeded.  */
 124         size /= 2;
 125
 126 #ifdef RLIMIT_RSS
 127         /* Leave a 1/16 margin for RSS to leave room for code, stack, etc.
 128            Exceeding RSS is not fatal, but can be quite slow.  */
 129         if (getrlimit(RLIMIT_RSS, &rlimit) == 0 && rlimit.rlim_cur / 16 * 15 < size)
 130                 size = rlimit.rlim_cur / 16 * 15;
 131 #endif
 132
 133         /* Use no less than the minimum. */
 134         return MAX (size, MIN_SPONGE_SIZE);
 135 }
 136
 137 void trapsignals (void) {
 138         ssize_t i = 0;
 139         static int const sig[] = {
 140                 /* The usual suspects.  */
 141                 SIGALRM, SIGHUP, SIGINT, SIGPIPE, SIGQUIT, SIGTERM,
 142 #ifdef SIGPOLL
 143                 SIGPOLL,
 144 #endif
 145 #ifdef SIGPROF
 146                 SIGPROF,
 147 #endif
 148 #ifdef SIGVTALRM
 149                 SIGVTALRM,
 150 #endif
 151 #ifdef SIGXCPU
 152                 SIGXCPU,
 153 #endif
 154 #ifdef SIGXFSZ
 155                 SIGXFSZ,
 156 #endif
 157         };
 158         int nsigs = sizeof(sig) / sizeof(sig[0]);
 159
 160 #if SA_NOCLDSTOP
 161         struct sigaction act;
 162
 163         sigemptyset(&caught_signals);
 164         for (i = 0; i < nsigs; i++) {
 165                 sigaction(sig[i], NULL, &act);
 166                 if (act.sa_handler != SIG_IGN)
 167                         sigaddset(&caught_signals, sig[i]);
 168                 }
 169
 170                 act.sa_handler = sighandler;
 171                 act.sa_mask = caught_signals;
 172                 act.sa_flags = 0;
 173
 174                 for (i = 0; i < nsigs; i++)
 175                         if (sigismember(&caught_signals, sig[i]))
 176                                 sigaction(sig[i], &act, NULL);
 177 #else
 178         for (i = 0; i < nsigs; i++)
 179                 if (signal(sig[i], SIG_IGN) != SIG_IGN) {
 180                         signal(sig[i], sighandler);
 181                         siginterrupt (sig[i], 1);
 182                 }
 183 #endif
 184 }
 185
 186 static void write_buff_tmp(char* buff, size_t length, FILE *fd) {
 187         if (fwrite(buff, length, 1, fd) < 1) {
 188                 perror("error writing buffer to temporary file");
 189                 fclose(fd);
 190                 exit(1);
 191         }
 192 }
 193
 194 static void write_buff_out(char* buff, size_t length, FILE *fd) {
 195         if (fwrite(buff, length, 1, fd) < 1) {
 196                 perror("error writing buffer to output file");
 197                 fclose(fd);
 198                 exit(1);
 199         }
 200 }
 201
 202 static void copy_tmpfile(FILE *tmpfile, FILE *outfd) {
 203         char buf[BUFF_SIZE];
 204         if (fseek(tmpfile, 0, SEEK_SET)) {
 205                 perror("could to seek to start of temporary file");
 206                 fclose(tmpfile);
 207                 exit(1);
 208         }
 209         while (fread(buf, BUFF_SIZE, 1, tmpfile) > 0) {
 210                 write_buff_out(buf, BUFF_SIZE, outfd);
 211         }
 212         if (ferror(tmpfile)) {
 213                 perror("read temporary file");
 214                 fclose(tmpfile);
 215                 exit(1);
 216         }
 217         fclose(tmpfile);
 218         fclose(outfd);
 219 }
 220
 221 int main (int argc, char **argv) {
 222         char *buf, *bufstart, *outname = NULL;
 223         size_t bufsize = BUFF_SIZE;
 224         size_t bufused = 0;
 225         FILE *tmpfile = 0;
 226         ssize_t i = 0;
 227         size_t mem_available = default_sponge_size();
 228
 229         if (argc > 2 || (argc == 2 && strcmp(argv[1], "-h") == 0)) {
 230                 usage();
 231         }
 232         bufstart = buf = malloc(bufsize);
 233         if (!buf) {
 234                 perror("failed to allocate memory");
 235                 exit(1);
 236         }
 237         while ((i = read(0, buf, bufsize - bufused)) > 0) {
 238                 bufused = bufused+i;
 239                 if (bufused == bufsize) {
 240                         if ((bufsize*2) >= mem_available) {
 241                                 if (!tmpfile) {
 242                                         /* umask(077); FIXME: Should we be setting umask, or using default?  */
 243                                         struct cs_status cs;
 244                                         int tmpfd;
 245
 246                                         trapsignals();
 247                                         cs = cs_enter();
 248                                         tmpfd = mkstemp(tmpname);
 249                                         atexit(onexit_cleanup); // solaris on_exit(onexit_cleanup, 0);
 250                                         cs_leave(cs);
 251                                         if (tmpfd < 0) {
 252                                                 perror("mkstemp failed");
 253                                                 exit(1);
 254                                         }
 255                                         tmpfile = fdopen(tmpfd, "w+");
 256                                 }
 257                                 write_buff_tmp(bufstart, bufused, tmpfile);
 258                                 bufused = 0;
 259                         }
 260                         else {
 261                                 bufsize *= 2;
 262                                 bufstart = realloc(bufstart, bufsize);
 263                                 if (!bufstart) {
 264                                         perror("failed to realloc memory");
 265                                         exit(1);
 266                                 }
 267                         }
 268                 }
 269                 buf = bufstart + bufused;
 270         }
 271         if (i < 0) {
 272                 perror("failed to read from stdin");
 273                 exit(1);
 274         }
 275         if (argc == 2) {
 276                 outname = argv[1];
 277         }
 278         if (tmpfile) {
 279                 /* write whatever we have in memory to tmpfile */
 280                 if (bufused)
 281                         write_buff_tmp(bufstart, bufused, tmpfile);
 282                 struct stat statbuf;
 283                 if (outname && !stat(outname, &statbuf)) {
 284                         /* regular file */
 285                         if (S_ISREG(statbuf.st_mode) && !fclose(tmpfile)) {
 286                                 if (rename(tmpname, outname)) {
 287                                         perror("error renaming temporary file to output file");
 288                                         exit(1);
 289                                 }
 290                         }
 291                         else {
 292                                 FILE *outfd = fopen(outname, "w");
 293                                 if (outfd < 0) {
 294                                         perror("error opening output file");
 295                                         exit(1);
 296                                 }
 297                                 copy_tmpfile(tmpfile, outfd);
 298                         }
 299                 }
 300                 else {
 301                         copy_tmpfile(tmpfile, stdout);
 302                 }
 303         }
 304         else {
 305                 FILE *outfd = stdout;
 306                 if (outname) {
 307                         outfd = fopen(outname, "w");
 308                         if (outfd < 0) {
 309                                 perror("error opening output file");
 310                                 exit(1);
 311                         }
 312                 }
 313                 if (bufused)
 314                         write_buff_out(bufstart, bufused, outfd);
 315                 fclose(outfd);
 316         }
 317
 318         return 0;
 319 }