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     sizeof(8192)
  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>: suck in 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         unlink(tmpname);
  78 }
  79
  80 static void onexit_cleanup (void) {
  81         struct cs_status cs = cs_enter();
  82         cleanup();
  83         cs_leave(cs);
  84 }
  85
  86 static void sighandler (int sig) {
  87         if (! SA_NOCLDSTOP)
  88                 signal(sig, SIG_IGN);
  89
  90         cleanup();
  91
  92         signal(sig, SIG_DFL);
  93         raise(sig);
  94 }
  95
  96 /* taken from coreutils sort */
  97 static size_t default_sponge_size (void) {
  98         /* Let MEM be available memory or 1/8 of total memory, whichever
  99            is greater.  */
 100         double avail = physmem_available();
 101         double total = physmem_total();
 102         double mem = MAX(avail, total / 8);
 103         struct rlimit rlimit;
 104
 105         /* Let SIZE be MEM, but no more than the maximum object size or
 106            system resource limits.  Avoid the MIN macro here, as it is not
 107            quite right when only one argument is floating point.  Don't
 108            bother to check for values like RLIM_INFINITY since in practice
 109            they are not much less than SIZE_MAX.  */
 110         size_t size = SIZE_MAX;
 111         if (mem < size)
 112                 size = mem;
 113         if (getrlimit(RLIMIT_DATA, &rlimit) == 0 && rlimit.rlim_cur < size)
 114                 size = rlimit.rlim_cur;
 115 #ifdef RLIMIT_AS
 116         if (getrlimit(RLIMIT_AS, &rlimit) == 0 && rlimit.rlim_cur < size)
 117                 size = rlimit.rlim_cur;
 118 #endif
 119
 120         /* Leave a large safety margin for the above limits, as failure can
 121            occur when they are exceeded.  */
 122         size /= 2;
 123
 124 #ifdef RLIMIT_RSS
 125         /* Leave a 1/16 margin for RSS to leave room for code, stack, etc.
 126            Exceeding RSS is not fatal, but can be quite slow.  */
 127         if (getrlimit(RLIMIT_RSS, &rlimit) == 0 && rlimit.rlim_cur / 16 * 15 < size)
 128                 size = rlimit.rlim_cur / 16 * 15;
 129 #endif
 130
 131         /* Use no less than the minimum. */
 132         return MAX (size, MIN_SPONGE_SIZE);
 133 }
 134
 135 void trapsignals (void) {
 136         ssize_t i = 0;
 137         static int const sig[] = {
 138                 /* The usual suspects.  */
 139                 SIGALRM, SIGHUP, SIGINT, SIGPIPE, SIGQUIT, SIGTERM,
 140 #ifdef SIGPOLL
 141                 SIGPOLL,
 142 #endif
 143 #ifdef SIGPROF
 144                 SIGPROF,
 145 #endif
 146 #ifdef SIGVTALRM
 147                 SIGVTALRM,
 148 #endif
 149 #ifdef SIGXCPU
 150                 SIGXCPU,
 151 #endif
 152 #ifdef SIGXFSZ
 153                 SIGXFSZ,
 154 #endif
 155         };
 156         int nsigs = sizeof(sig) / sizeof(sig[0]);
 157
 158 #if SA_NOCLDSTOP
 159         struct sigaction act;
 160
 161         sigemptyset(&caught_signals);
 162         for (i = 0; i < nsigs; i++) {
 163                 sigaction(sig[i], NULL, &act);
 164                 if (act.sa_handler != SIG_IGN)
 165                         sigaddset(&caught_signals, sig[i]);
 166                 }
 167
 168                 act.sa_handler = sighandler;
 169                 act.sa_mask = caught_signals;
 170                 act.sa_flags = 0;
 171
 172                 for (i = 0; i < nsigs; i++)
 173                         if (sigismember(&caught_signals, sig[i]))
 174                                 sigaction(sig[i], &act, NULL);
 175 #else
 176         for (i = 0; i < nsigs; i++)
 177                 if (signal(sig[i], SIG_IGN) != SIG_IGN) {
 178                         signal(sig[i], sighandler);
 179                         siginterrupt (sig[i], 1);
 180                 }
 181 #endif
 182 }
 183
 184 int main (int argc, char **argv) {
 185         char *buf, *bufstart, *outname = NULL;
 186         size_t bufsize = BUFF_SIZE;
 187         size_t bufused = 0;
 188         FILE *tmpfile = 0;
 189         ssize_t i = 0;
 190         size_t mem_available = default_sponge_size();
 191
 192         if (argc > 2 || (argc == 2 && strcmp(argv[1], "-h") == 0)) {
 193                 usage();
 194         }
 195         bufstart = buf = malloc(bufsize);
 196         if (!buf) {
 197                 perror("failed to allocate memory");
 198                 exit(1);
 199         }
 200
 201         trapsignals();
 202
 203         while ((i = read(0, buf, bufsize - bufused)) > 0) {
 204                 bufused = bufused+i;
 205                 if (bufused == bufsize) {
 206                         if (bufsize >= mem_available) {
 207                                 if (!tmpfile) {
 208                                         umask(077);
 209                                         struct cs_status cs = cs_enter();
 210                                         int tmpfd = mkstemp(tmpname);
 211                                         atexit(onexit_cleanup); // if solaris on_exit(onexit_cleanup, 0);
 212                                         cs_leave(cs);
 213                                         if (tmpfd < 0) {
 214                                                 perror("mkstemp");
 215                                                 exit(1);
 216                                         }
 217                                         tmpfile = fdopen(tmpfd, "w+");
 218                                 }
 219                                 if (fwrite(bufstart, bufsize, 1, tmpfile) < 1) {
 220                                         perror("writing to tempory file failed");
 221                                         fclose(tmpfile);
 222                                         exit(1);
 223                                 }
 224                                 bufused = 0;
 225                         }
 226                         else {
 227                                 bufsize *= 2;
 228                                 bufstart = realloc(bufstart, bufsize);
 229                                 if (!bufstart) {
 230                                         perror("failed to realloc memory");
 231                                         exit(1);
 232                                 }
 233                         }
 234                 }
 235                 buf = bufstart + bufused;
 236         }
 237         if (i < 0) {
 238                 perror("failed to read from stdin");
 239                 exit(1);
 240         }
 241         if (argc == 2) {
 242                 outname = argv[1];
 243         }
 244         if (tmpfile) {
 245                 if (fwrite(bufstart, bufused, 1, tmpfile) < 1) {
 246                         perror("write tmpfile");
 247                         fclose(tmpfile);
 248                         exit(1);
 249                 }
 250                 if (outname) {
 251                         fclose(tmpfile);
 252                         if (rename(tmpname, outname)) {
 253                                 perror("error renaming temporary file to output file");
 254                                 exit(1);
 255                         }
 256                 }
 257                 else {
 258                         if (fseek(tmpfile, 0, SEEK_SET)) {
 259                                 perror("could to seek to start of temporary file");
 260                                 fclose(tmpfile);
 261                                 exit(1);
 262                         }
 263                         while (fread( buf, BUFF_SIZE, 1, tmpfile) < 1) {
 264                                 if (fwrite(buf, BUFF_SIZE, 1, stdout) < 1) {
 265                                         perror("error writing out merged file");
 266                                         exit(1);
 267                                 }
 268                         }
 269                         fclose(tmpfile);
 270                         unlink(tmpname);
 271                 }
 272         }
 273         else {
 274                 if (outname) {
 275                         FILE *outfd = fopen(outname, "w");
 276                         if (outfd < 0) {
 277                                 perror("error opening output file");
 278                                 exit(1);
 279                         }
 280                         if (fwrite(bufstart, bufused, 1, outfd) < 1) {
 281                                 perror("error writing out merged file");
 282                                 exit(1);
 283                         }
 284                 }
 285                 else if (fwrite(bufstart, bufused, 1, stdout) < 1) {
 286                         perror("error writing out merged file");
 287                         exit(1);
 288                 }
 289         }
 290         return 0;
 291 }