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 MIN_SPONGE_SIZE     BUFF_SIZE
  41 #define BUFF_SIZE           8192
  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         // XXX I'd catch signals or writes errors here, but I
 210         // I don't think it matters as the file is overwritten
 211         while(fread(buf, BUFF_SIZE, 1, tmpfile) == 1) {
 212                 write_buff_out(buf, BUFF_SIZE, outfd);
 213         }
 214         fclose(tmpfile);
 215         fclose(outfd);
 216 }
 217
 218 int main (int argc, char **argv) {
 219         char *buf, *bufstart, *outname = NULL;
 220         size_t bufsize = BUFF_SIZE;
 221         size_t bufused = 0;
 222         FILE *tmpfile = 0;
 223         ssize_t i = 0;
 224         size_t mem_available = default_sponge_size();
 225
 226         if (argc > 2 || (argc == 2 && strcmp(argv[1], "-h") == 0)) {
 227                 usage();
 228         }
 229         bufstart = buf = malloc(bufsize);
 230         if (!buf) {
 231                 perror("failed to allocate memory");
 232                 exit(1);
 233         }
 234         while ((i = read(0, buf, bufsize - bufused)) > 0) {
 235                 bufused = bufused+i;
 236                 if (bufused == bufsize) {
 237                         if ((bufsize*2) >= mem_available) {
 238                                 if (!tmpfile) {
 239                                         /* umask(077); FIXME: Should we be setting umask, or using default?  */
 240                                         struct cs_status cs = cs_enter();
 241                                         int tmpfd = mkstemp(tmpname);
 242                                         atexit(onexit_cleanup); // solaris on_exit(onexit_cleanup, 0);
 243                                         trapsignals();
 244                                         cs_leave(cs);
 245                                         if (tmpfd < 0) {
 246                                                 perror("mkstemp failed");
 247                                                 exit(1);
 248                                         }
 249                                         tmpfile = fdopen(tmpfd, "w+");
 250                                 }
 251                                 write_buff_tmp(bufstart, bufused, tmpfile);
 252                                 bufused = 0;
 253                         }
 254                         else {
 255                                 bufsize *= 2;
 256                                 bufstart = realloc(bufstart, bufsize);
 257                                 if (!bufstart) {
 258                                         perror("failed to realloc memory");
 259                                         exit(1);
 260                                 }
 261                         }
 262                 }
 263                 buf = bufstart + bufused;
 264         }
 265         if (i < 0) {
 266                 perror("failed to read from stdin");
 267                 exit(1);
 268         }
 269         if (argc == 2) {
 270                 outname = argv[1];
 271         }
 272         if (tmpfile) {
 273                 /* write whatever we have in memory to tmpfile */
 274                 write_buff_tmp(bufstart, bufused, tmpfile);
 275                 struct stat statbuf;
 276                 if (outname && !stat(outname, &statbuf)) {
 277                         /* regular file */
 278                         if (S_ISREG(statbuf.st_mode) && !fclose(tmpfile)) {
 279                                 if (rename(tmpname, outname)) {
 280                                         perror("error renaming temporary file to output file");
 281                                         exit(1);
 282                                 }
 283                         }
 284                         else {
 285                                 FILE *outfd = fopen(outname, "w");
 286                                 if (outfd < 0) {
 287                                         perror("error opening output file");
 288                                         exit(1);
 289                                 }
 290                                 copy_tmpfile(tmpfile, outfd);
 291                         }
 292                 }
 293                 else {
 294                         copy_tmpfile(tmpfile, stdout);
 295                 }
 296         }
 297         else {
 298                 FILE *outfd = stdout;
 299                 if (outname) {
 300                         outfd = fopen(outname, "w");
 301                         if (outfd < 0) {
 302                                 perror("error opening output file");
 303                                 exit(1);
 304                         }
 305                 }
 306                 write_buff_out(bufstart, bufused, outfd);
 307                 fclose(outfd);
 308         }
 309
 310         return 0;
 311 }