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 char *tmpname = NULL;
  43
  44 void usage() {
  45         printf("sponge <file>: soak up all input from stdin and write it to <file>\n");
  46         exit(0);
  47 }
  48
  49 /* all the signal stuff copied from gnu sort */
  50
  51 /* The set of signals that are caught.  */
  52 static sigset_t caught_signals;
  53
  54 /* Critical section status.  */
  55 struct cs_status {
  56         int valid; // was bool
  57         sigset_t sigs;
  58 };
  59
  60 /* Enter a critical section.  */
  61 static struct cs_status cs_enter (void) {
  62         struct cs_status status;
  63         status.valid = (sigprocmask(SIG_BLOCK, &caught_signals, &status.sigs) == 0);
  64         return status;
  65 }
  66
  67 /* Leave a critical section.  */
  68 static void cs_leave (struct cs_status status) {
  69         if (status.valid) {
  70                 /* Ignore failure when restoring the signal mask. */
  71                 sigprocmask(SIG_SETMASK, &status.sigs, NULL);
  72         }
  73 }
  74
  75 static void cleanup () {
  76         if (tmpname) {
  77                 unlink(tmpname);
  78         }
  79 }
  80
  81 static void onexit_cleanup (void) {
  82         struct cs_status cs = cs_enter();
  83         cleanup();
  84         cs_leave(cs);
  85 }
  86
  87 static void sighandler (int sig) {
  88         if (! SA_NOCLDSTOP)
  89                 signal(sig, SIG_IGN);
  90
  91         cleanup();
  92
  93         signal(sig, SIG_DFL);
  94         raise(sig);
  95 }
  96
  97 /* taken from coreutils sort */
  98 static size_t default_sponge_size (void) {
  99         /* Let MEM be available memory or 1/8 of total memory, whichever
 100            is greater.  */
 101         double avail = physmem_available();
 102         double total = physmem_total();
 103         double mem = MAX(avail, total / 8);
 104         struct rlimit rlimit;
 105
 106         /* Let SIZE be MEM, but no more than the maximum object size or
 107            system resource limits.  Avoid the MIN macro here, as it is not
 108            quite right when only one argument is floating point.  Don't
 109            bother to check for values like RLIM_INFINITY since in practice
 110            they are not much less than SIZE_MAX.  */
 111         size_t size = SIZE_MAX;
 112         if (mem < size)
 113                 size = mem;
 114         if (getrlimit(RLIMIT_DATA, &rlimit) == 0 && rlimit.rlim_cur < size)
 115                 size = rlimit.rlim_cur;
 116 #ifdef RLIMIT_AS
 117         if (getrlimit(RLIMIT_AS, &rlimit) == 0 && rlimit.rlim_cur < size)
 118                 size = rlimit.rlim_cur;
 119 #endif
 120
 121         /* Leave a large safety margin for the above limits, as failure can
 122            occur when they are exceeded.  */
 123         size /= 2;
 124
 125 #ifdef RLIMIT_RSS
 126         /* Leave a 1/16 margin for RSS to leave room for code, stack, etc.
 127            Exceeding RSS is not fatal, but can be quite slow.  */
 128         if (getrlimit(RLIMIT_RSS, &rlimit) == 0 && rlimit.rlim_cur / 16 * 15 < size)
 129                 size = rlimit.rlim_cur / 16 * 15;
 130 #endif
 131
 132         /* Use no less than the minimum. */
 133         return MAX (size, MIN_SPONGE_SIZE);
 134 }
 135
 136 void trapsignals (void) {
 137         ssize_t i = 0;
 138         static int const sig[] = {
 139                 /* The usual suspects.  */
 140                 SIGALRM, SIGHUP, SIGINT, SIGPIPE, SIGQUIT, SIGTERM,
 141 #ifdef SIGPOLL
 142                 SIGPOLL,
 143 #endif
 144 #ifdef SIGPROF
 145                 SIGPROF,
 146 #endif
 147 #ifdef SIGVTALRM
 148                 SIGVTALRM,
 149 #endif
 150 #ifdef SIGXCPU
 151                 SIGXCPU,
 152 #endif
 153 #ifdef SIGXFSZ
 154                 SIGXFSZ,
 155 #endif
 156         };
 157         int nsigs = sizeof(sig) / sizeof(sig[0]);
 158
 159 #if SA_NOCLDSTOP
 160         struct sigaction act;
 161
 162         sigemptyset(&caught_signals);
 163         for (i = 0; i < nsigs; i++) {
 164                 sigaction(sig[i], NULL, &act);
 165                 if (act.sa_handler != SIG_IGN)
 166                         sigaddset(&caught_signals, sig[i]);
 167                 }
 168
 169                 act.sa_handler = sighandler;
 170                 act.sa_mask = caught_signals;
 171                 act.sa_flags = 0;
 172
 173                 for (i = 0; i < nsigs; i++)
 174                         if (sigismember(&caught_signals, sig[i]))
 175                                 sigaction(sig[i], &act, NULL);
 176 #else
 177         for (i = 0; i < nsigs; i++)
 178                 if (signal(sig[i], SIG_IGN) != SIG_IGN) {
 179                         signal(sig[i], sighandler);
 180                         siginterrupt (sig[i], 1);
 181                 }
 182 #endif
 183 }
 184
 185 static void write_buff_tmp(char* buff, size_t length, FILE *fd) {
 186         if (fwrite(buff, length, 1, fd) < 1) {
 187                 perror("error writing buffer to temporary file");
 188                 fclose(fd);
 189                 exit(1);
 190         }
 191 }
 192
 193 static void write_buff_tmp_finish (char* buff, size_t length, FILE *fd) {
 194         if (length)
 195                 write_buff_tmp(buff, length, fd);
 196         if (fflush(fd) != 0) {
 197                 perror("fflush");
 198                 exit(1);
 199         }
 200 }
 201
 202 static void write_buff_out (char* buff, size_t length, FILE *fd) {
 203         if (fwrite(buff, length, 1, fd) < 1) {
 204                 perror("error writing buffer to output file");
 205                 fclose(fd);
 206                 exit(1);
 207         }
 208 }
 209
 210 static void copy_tmpfile (FILE *tmpfile, FILE *outfile, char *buf, size_t size) {
 211         ssize_t i;
 212         if (lseek(fileno(tmpfile), 0, SEEK_SET)) {
 213                 perror("could to seek to start of temporary file");
 214                 fclose(tmpfile);
 215                 exit(1);
 216         }
 217         while ((i = read(fileno(tmpfile), buf, size)) > 0) {
 218                 write_buff_out(buf, i, outfile);
 219         }
 220         if (i == -1) {
 221                 perror("read temporary file");
 222                 fclose(tmpfile);
 223                 exit(1);
 224         }
 225         if (fclose(tmpfile) != 0) {
 226                 perror("read temporary file");
 227                 exit(1);
 228         }
 229         if (fclose(outfile) != 0) {
 230                 perror("error writing buffer to output file");
 231                 exit(1);
 232         }
 233 }
 234
 235 FILE *open_tmpfile (void) {
 236         struct cs_status cs;
 237         int tmpfd;
 238         FILE *tmpfile;
 239         mode_t mask;
 240         char *tmpdir;
 241         char const * const template="%s/sponge.XXXXXX";
 242
 243         trapsignals();
 244         cs = cs_enter();
 245         tmpdir = getenv("TMPDIR");
 246         if (tmpdir == NULL)
 247                 tmpdir = "/tmp";
 248         /* Subtract 2 for `%s' and add 1 for the trailing NULL. */
 249         tmpname=malloc(strlen(tmpdir) + strlen(template) - 2 + 1);
 250         if (! tmpname) {
 251                 perror("failed to allocate memory");
 252                 exit(1);
 253         }
 254         sprintf(tmpname, template, tmpdir);
 255         mask=umask(077);
 256         tmpfd = mkstemp(tmpname);
 257         umask(mask);
 258         atexit(onexit_cleanup); // solaris on_exit(onexit_cleanup, 0);
 259         cs_leave(cs);
 260
 261         if (tmpfd < 0) {
 262                 perror("mkstemp failed");
 263                 exit(1);
 264         }
 265         tmpfile = fdopen(tmpfd, "w+");
 266         if (! tmpfile) {
 267                 perror("fdopen");
 268                 exit(1);
 269         }
 270         return tmpfile;
 271 }
 272
 273 int main (int argc, char **argv) {
 274         char *buf, *bufstart, *outname = NULL;
 275         size_t bufsize = BUFF_SIZE;
 276         size_t bufused = 0;
 277         FILE *outfile, *tmpfile = 0;
 278         ssize_t i = 0;
 279         size_t mem_available = default_sponge_size();
 280         int tmpfile_used=0;
 281
 282         if (argc > 2 || (argc == 2 && strcmp(argv[1], "-h") == 0)) {
 283                 usage();
 284         }
 285         if (argc == 2) {
 286                 outname = argv[1];
 287         }
 288
 289         tmpfile = open_tmpfile();
 290         bufstart = buf = malloc(bufsize);
 291         if (!buf) {
 292                 perror("failed to allocate memory");
 293                 exit(1);
 294         }
 295         while ((i = read(0, buf, bufsize - bufused)) > 0) {
 296                 bufused = bufused+i;
 297                 if (bufused == bufsize) {
 298                         if ((bufsize*2) >= mem_available) {
 299                                 write_buff_tmp(bufstart, bufused, tmpfile);
 300                                 bufused = 0;
 301                                 tmpfile_used = 1;
 302                         }
 303                         else {
 304                                 bufsize *= 2;
 305                                 bufstart = realloc(bufstart, bufsize);
 306                                 if (!bufstart) {
 307                                         perror("failed to realloc memory");
 308                                         exit(1);
 309                                 }
 310                         }
 311                 }
 312                 buf = bufstart + bufused;
 313         }
 314         if (i < 0) {
 315                 perror("failed to read from stdin");
 316                 exit(1);
 317         }
 318
 319         if (outname) {
 320                 mode_t mode;
 321                 struct stat statbuf;
 322                 int exists = (lstat(outname, &statbuf) == 0);
 323
 324                 write_buff_tmp_finish(bufstart, bufused, tmpfile);
 325
 326                 /* Set temp file mode to match either
 327                  * the old file mode, or the default file
 328                  * mode for a newly created file. */
 329                 if (exists) {
 330                         mode = statbuf.st_mode;
 331                 }
 332                 else {
 333                         mode_t mask = umask(0);
 334                         umask(mask);
 335                         mode = 0666 & ~mask;
 336                 }
 337                 if (chmod(tmpname, mode) != 0) {
 338                         perror("chmod");
 339                         exit(1);
 340                 }
 341
 342                 /* If it's a regular file, or does not yet exist,
 343                  * attempt a fast rename of the temp file. */
 344                 if (((exists &&
 345                       S_ISREG(statbuf.st_mode) &&
 346                       ! S_ISLNK(statbuf.st_mode)
 347                      ) || ! exists) &&
 348                     rename(tmpname, outname) == 0) {
 349                         tmpname=NULL; /* don't try to cleanup tmpname */
 350                 }
 351                 else {
 352                         /* Fall back to slow copy. */
 353                         outfile = fopen(outname, "w");
 354                         if (!outfile) {
 355                                 perror("error opening output file");
 356                                 exit(1);
 357                         }
 358                         copy_tmpfile(tmpfile, outfile, bufstart, bufsize);
 359                 }
 360         }
 361         else {
 362                 if (tmpfile_used) {
 363                         write_buff_tmp_finish(bufstart, bufused, tmpfile);
 364                         copy_tmpfile(tmpfile, stdout, bufstart, bufsize);
 365                 }
 366                 else if (bufused) {
 367                         /* buffer direct to stdout, no tmpfile */
 368                         write_buff_out(bufstart, bufused, stdout);
 369                 }
 370         }
 371
 372         return 0;
 373 }